Update app.py

app.py

@@ -1,7 +1,6 @@
+
 import os
 import torch
-import cv2
-import numpy as np
 import gradio as gr
 from PIL import Image
 
@@ -20,39 +19,14 @@ else:
     print(f"Model weights found at: {best_path}")
     model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
 
-# 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']
-    weapons = [cls for cls in detected_classes if cls in WEAPON_CLASSES]
-    threat_message = "Threat detected: Be careful" if weapons else "No threat detected"
+    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])
 
-# 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
+# Gradio interface
 iface = gr.Interface(
     fn=detect_weapons,
     inputs=gr.Image(type="numpy", label="Upload Image"),
@@ -61,21 +35,7 @@ iface = gr.Interface(
         gr.Image(label="Detected Image"),
     ],
     title="Weapon Detection AI",
-    description="Upload an image to detect weapons like bombs, guns, and pistols. For real-time detection, use the live camera option."
+    description="Upload an image to detect weapons like bombs, guns, and pistols."
 )
 
-# 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()
+iface.launch()