Update app.py

app.py

@@ -3,26 +3,50 @@ import torch
 import gradio as gr
 from PIL import Image
 
-# Step 1: Load the custom YOLOv5 model
-# Make sure the path to the model is correct, either locally or from Hugging Face model hub
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')  # Adjust the path if needed
+# Step 1: Search for best.pt in the training directory
+base_path = "yolov5/runs/train/"
+best_path = None
 
-# Step 2: Define weapon classes to detect
-weapon_classes = ['pistol', 'gun', 'rifle', 'shotgun', 'handgun']
+# Search through the directory structure to find best.pt
+for root, dirs, files in os.walk(base_path):
+    if "best.pt" in files:
+        best_path = os.path.join(root, "best.pt")
+        break
+
+# Step 2: If best.pt is not found, use pre-trained weights
+if best_path is None:
+    print("Trained weights (best.pt) not found.")
+    print("Using pre-trained YOLOv5 weights (yolov5s.pt) instead.")
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Load pre-trained weights
+else:
+    print(f"Model weights found at: {best_path}")
+    model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
+
+# Step 3: Define weapon classes to detect
+weapon_classes = ['bomb', 'gun', 'pistol', 'Automatic', 'Rifle', 'Bazooka', 
+                  'Handgun', 'Knife', 'Grenade Launcher', 'Shotgun', 'SMG', 
+                  'Sniper', 'Sword']  # Adjust based on your dataset
 
 def detect_weapons(image):
     results = model(image)
-    
+
+    # Print available model class names to check for class mismatches
+    model_classes = results.names  # This should give the list of class labels used by the model
+    print("Model class names:", model_classes)
+
     # Filter detections by confidence threshold (0.5 or higher)
     confidence_threshold = 0.5
     filtered_results = results.pandas().xyxy[0][results.pandas().xyxy[0]['confidence'] >= confidence_threshold]
-    
+
     # Get the detected classes with high confidence
     detected_classes = filtered_results['name'].unique()
 
+    # Print detected classes for debugging
+    print("Detected classes:", detected_classes)
+
     # Check if any of the detected objects are weapons
     detected_threats = [weapon for weapon in weapon_classes if weapon in detected_classes]
-    
+
     # Determine threat message based on weapons detected
     if detected_threats:
         threat_message = "Threat detected: Be careful"
@@ -35,7 +59,7 @@ def detect_weapons(image):
     # Render the image with bounding boxes
     return f"{threat_message}\nDetected objects: {detected_objects}", Image.fromarray(results.render()[0])
 
-# Step 3: Gradio Interface
+# Step 4: Gradio Interface
 def inference(image):
     threat, detected_image = detect_weapons(image)
     return threat, detected_image
@@ -48,8 +72,8 @@ iface = gr.Interface(
         gr.Image(label="Detected Image"),
     ],
     title="Weapon Detection AI",
-    description="Upload an image to detect weapons like pistols, guns, and rifles."
+    description="Upload an image to detect weapons like bombs, guns, and pistols."
 )
 
-# Step 4: Launch Gradio App
+# Step 5: Launch Gradio App
 iface.launch()