Update app.py

app.py

@@ -1,41 +1,55 @@
-
 import os
 import torch
 import gradio as gr
 from PIL import Image
 
-# Load YOLOv5 model
-base_path = "yolov5/runs/train/"
-best_path = None
-for root, dirs, files in os.walk(base_path):
-    if "best.pt" in files:
-        best_path = os.path.join(root, "best.pt")
-        break
-
-if best_path is None:
-    print("Trained weights (best.pt) not found. Using pre-trained weights.")
-    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
-else:
-    print(f"Model weights found at: {best_path}")
-    model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
-
-# Detection function
+# Step 1: Load the custom YOLOv5 model
+# Adjust the model path based on your deployment setup (Hugging Face model hub or local path)
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')  # If you use Hugging Face model hub, use the correct identifier
+
+# Step 2: Define weapon classes to detect
+weapon_classes = ['pistol', 'gun', 'rifle', 'shotgun', 'handgun']
+
 def detect_weapons(image):
     results = model(image)
-    detected_classes = results.pandas().xyxy[0]['name'].unique()
-    threat_message = "Threat detected: Be careful" if len(detected_classes) > 0 else "No threat detected"
-    return threat_message, Image.fromarray(results.render()[0])
+    
+    # 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()
+
+    # 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"
+    else:
+        threat_message = "No threat detected. But all other features are good."
+
+    # Create a string with the detected objects' names
+    detected_objects = ', '.join(detected_classes)
+
+    # Render the image with bounding boxes
+    return f"{threat_message}\nDetected objects: {detected_objects}", Image.fromarray(results.render()[0])
+
+# Step 3: Gradio Interface
+def inference(image):
+    threat, detected_image = detect_weapons(image)
+    return threat, detected_image
 
-# Gradio interface
 iface = gr.Interface(
-    fn=detect_weapons,
+    fn=inference,
     inputs=gr.Image(type="numpy", label="Upload Image"),
     outputs=[
         gr.Textbox(label="Threat Detection"),
         gr.Image(label="Detected Image"),
     ],
     title="Weapon Detection AI",
-    description="Upload an image to detect weapons like bombs, guns, and pistols."
+    description="Upload an image to detect weapons like pistols, guns, and rifles."
 )
 
-iface.launch()
+# Step 4: Launch Gradio App
+iface.launch()