Update app.py

app.py

@@ -1,56 +1,81 @@
 import os
 import torch
+import cv2
+import numpy as np
 import gradio as gr
 from PIL import Image
 
-# Step 1: Search for best.pt in the training directory
+# Load YOLOv5 model
 base_path = "yolov5/runs/train/"
 best_path = None
-
-# 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.")
-
-    # Fallback to a pre-trained model if best.pt is not available
-    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Load pre-trained weights
-
+    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}")
-    # Load YOLOv5 model with the correct path
     model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
 
-# Step 3: Detection Function
+# Define weapon-related classes
+WEAPON_CLASSES = ["pistol", "rifle", "knife", "bomb", "gun", "weapon"]
+
+# Detection function
 def detect_weapons(image):
     results = model(image)
-    detected_classes = results.pandas().xyxy[0]['name'].unique()
-    
-    # Check for threats
-    threat_message = "Threat detected: Be careful" if len(detected_classes) > 0 else "No threat detected"
+    detected_classes = results.pandas().xyxy[0]['name']
+    weapons = [cls for cls in detected_classes if cls in WEAPON_CLASSES]
+    threat_message = "Threat detected: Be careful" if weapons else "No threat detected"
     return threat_message, Image.fromarray(results.render()[0])
 
-# Step 4: Gradio Interface
-def inference(image):
-    threat, detected_image = detect_weapons(image)
-    return threat, detected_image
+# Real-time detection from webcam
+def detect_from_camera():
+    cap = cv2.VideoCapture(0)  # Open webcam (use 0 for default camera)
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        # Convert frame (OpenCV format) to PIL image
+        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        results = model(image)
+        detected_classes = results.pandas().xyxy[0]['name']
+        weapons = [cls for cls in detected_classes if cls in WEAPON_CLASSES]
+        annotated_frame = np.array(results.render()[0])
+        
+        # Display live feed with annotations
+        cv2.imshow("Weapon Detection", cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR))
+        if cv2.waitKey(1) & 0xFF == ord('q'):  # Press 'q' to quit
+            break
+    cap.release()
+    cv2.destroyAllWindows()
 
+# Gradio interface for image uploads
 iface = gr.Interface(
-    fn=inference,
+    fn=detect_weapons,
     inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=[ 
-        gr.Textbox(label="Threat Detection"), 
-        gr.Image(label="Detected 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 bombs, guns, and pistols. For real-time detection, use the live camera option."
 )
 
-# Step 5: Launch Gradio App
-iface.launch()
+# Main entry point
+def main():
+    print("Select mode:")
+    print("1. Image Upload (Gradio Interface)")
+    print("2. Live Camera Detection")
+    choice = input("Enter your choice (1 or 2): ").strip()
+    if choice == "1":
+        iface.launch()
+    elif choice == "2":
+        detect_from_camera()
+    else:
+        print("Invalid choice. Please restart and select 1 or 2.")
 
+if __name__ == "__main__":
+    main()