Update app.py

app.py

@@ -4,51 +4,50 @@ import streamlit as st
 from datetime import datetime
 from tempfile import NamedTemporaryFile
 from ultralytics import YOLO
-import os
-os.system("pip install opencv-python")
-import cv2
 
-
-# Initialize YOLOv8 model
-model = YOLO("best.pt")  # Replace with your YOLOv8 model path
+# Initialize YOLO models
+weapon_model = YOLO("best.pt")  # Weapon detection model
+effect_model = YOLO("yolov8n.pt")  # Pretrained YOLO model
+
+# Load class labels for effect_model (YOLOv8n)
+coco_labels = [
+    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
+    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
+    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
+    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
+    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
+    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
+    "potted plant", "bed", "dining table", "toilet", "TV", "laptop", "mouse", "remote", "keyboard",
+    "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
+    "scissors", "teddy bear", "hair drier", "toothbrush"
+]
 
 # Output folder for storing detected clips
 output_folder = "detected_clips"
 os.makedirs(output_folder, exist_ok=True)
 
-# Initialize session state variables
+# Streamlit UI
+st.title("Weapon Detection System with Object Recognition")
+frame_placeholder = st.empty()
+log_placeholder = st.empty()
+
+# Session state variables
 if "log" not in st.session_state:
     st.session_state.log = ""
 if "stop_camera" not in st.session_state:
     st.session_state.stop_camera = False
 
-# Streamlit UI
-st.title("Weapon Detection System")
-st.write("This is a weapon detection system using YOLOv8 and OpenCV.")
-
-# Placeholder for video feed (Camera Stream)
-frame_placeholder = st.empty()
-
-# Placeholder for logs (⚠ Now defined before use)
-log_placeholder = st.empty()
-
-# Function to write log messages in real-time
 def write_log(message):
     st.session_state.log += message + "\n"
-    # Display logs in a scrolling text area
     log_placeholder.text_area("Detection Log", value=st.session_state.log, height=300, max_chars=5000, disabled=True)
 
-# Function to process video feed and detect weapons
 def process_video_feed(video_file=None):
     st.session_state.stop_camera = False
 
-    # Open webcam or video file
     if video_file:
-        # Save the uploaded video to a temporary file
         with NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
             temp_file.write(video_file.read())
             temp_file_path = temp_file.name
-
         cap = cv2.VideoCapture(temp_file_path)
     else:
         cap = cv2.VideoCapture(0)
@@ -67,24 +66,51 @@ def process_video_feed(video_file=None):
         if not ret:
             break
 
-        # Run YOLOv8 inference
-        results = model(frame, verbose=False)
-        annotated_frame = results[0].plot()
+        # Run YOLO inference on both models
+        weapon_results = weapon_model(frame, verbose=False)
+        effect_results = effect_model(frame, verbose=False)
 
+        annotated_frame = frame.copy()
         detected = False
         weapon_name = ""
         now = datetime.now()
 
-        # Check for weapon detection
-        for result in results[0].boxes:
-            class_id = int(result.cls[0])
-            confidence = float(result.conf[0])
-            if class_id == 0 and confidence > 0.5:
-                detected = True
-                weapon_name = "Weapon"
-                break
-
-        # If weapon detected, start saving the video
+        # 🟥 **Weapon Model Bounding Boxes**
+        for result in weapon_results:
+            if result.boxes is not None and len(result.boxes) > 0:
+                for box in result.boxes:
+                    class_id = int(box.cls.item())
+                    confidence = float(box.conf.item())
+
+                    if class_id == 0 and confidence > 0.5:  # Assuming class 0 is weapon
+                        detected = True
+                        weapon_name = "Weapon"
+
+                        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                        cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 0, 255), 2)  # Red
+                        cv2.putText(
+                            annotated_frame, f"{weapon_name} ({confidence:.2f})", 
+                            (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2
+                        )
+
+        # 🔵 **Effect Model Bounding Boxes**
+        for result in effect_results:
+            if result.boxes is not None and len(result.boxes) > 0:
+                for box in result.boxes:
+                    class_id = int(box.cls.item())
+                    confidence = float(box.conf.item())
+
+                    if confidence > 0.5:  # General detection for all objects
+                        object_name = coco_labels[class_id] if class_id < len(coco_labels) else f"Object-{class_id}"
+
+                        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                        cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (255, 0, 0), 2)  # Blue
+                        cv2.putText(
+                            annotated_frame, f"{object_name} ({confidence:.2f})", 
+                            (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2
+                        )
+
+        # **Start Recording if Weapon is Detected**
         if detected:
             if not recording:
                 formatted_date = now.strftime("%d-%m-%y")
@@ -102,8 +128,6 @@ def process_video_feed(video_file=None):
                 write_log(f"🎥 Recording started for {output_video_path}")
             
             out.write(frame)
-
-            # Add detection log
             timestamp = now.strftime("%d/%m/%y %I:%M:%S %p")
             write_log(f"⚠️ {weapon_name} detected at {timestamp}")