Update app.py

app.py

@@ -5,74 +5,86 @@ from ultralytics import YOLO
 import cv2
 import gradio as gr
 import numpy as np
+import torch
 from ultralytics import YOLO
 
-# Load the YOLOv8 model
-model = YOLO("yolov8n.pt")  # Using pre-trained YOLOv8 nano model
-
-# Object classes in YOLOv8
-CLASS_NAMES = model.names
-HUMAN_CLASS_ID = 0  # Class ID for "person" in YOLO
+# Load YOLOv8 model (pre-trained)
+model = YOLO("yolov8n.pt")  # Use a small YOLOv8 model for efficiency
 
+# Video Inference Function
 def vid_inf(vid_path, contour_thresh):
     cap = cv2.VideoCapture(vid_path)
-
-    # Get the video frames' width and height
+    
+    # Get video properties
     frame_width = int(cap.get(3))
     frame_height = int(cap.get(4))
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     frame_size = (frame_width, frame_height)
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     output_video = "output_recorded.mp4"
+    
+    # Video Writer
     out = cv2.VideoWriter(output_video, fourcc, fps, frame_size)
 
+    # Background Subtraction Model
     backSub = cv2.createBackgroundSubtractorMOG2(history=200, varThreshold=25, detectShadows=True)
-
+    
+    # Check if video opened successfully
     if not cap.isOpened():
         print("Error opening video file")
-        return
+        return None, None
 
     count = 0
+
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         
-        # YOLOv8 Object Detection
-        results = model(frame)
-        detected_boxes = []
-        
-        for result in results:
-            for box in result.boxes:
-                class_id = int(box.cls[0].item())
-                conf = box.conf[0].item()
-                
-                if class_id != HUMAN_CLASS_ID and conf > 0.5:  # Ignore humans, detect other objects
-                    x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates
-                    detected_boxes.append((x1, y1, x2, y2))
-
+        # Apply Background Subtraction (Motion Detection)
         fg_mask = backSub.apply(frame)
         retval, mask_thresh = cv2.threshold(fg_mask, 200, 255, cv2.THRESH_BINARY)
+
+        # Erosion to remove noise
         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
         mask_eroded = cv2.morphologyEx(mask_thresh, cv2.MORPH_OPEN, kernel)
 
+        # Find Contours (Potential Moving Objects)
         contours, _ = cv2.findContours(mask_eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-        min_contour_area = contour_thresh
-        large_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > min_contour_area]
+        large_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > contour_thresh]
 
         frame_out = frame.copy()
 
-        # Draw bounding boxes only on non-human moving objects
-        for cnt in large_contours:
-            x, y, w, h = cv2.boundingRect(cnt)
-            for (x1, y1, x2, y2) in detected_boxes:
-                if x > x1 and y > y1 and (x + w) < x2 and (y + h) < y2:  # Ensure it's inside an object
-                    frame_out = cv2.rectangle(frame_out, (x, y), (x + w, y + h), (0, 0, 200), 3)
+        # Run YOLOv8 Object Detection
+        results = model(frame, verbose=False)  # Perform object detection
+
+        for result in results:
+            boxes = result.boxes.xyxy  # Bounding boxes
+            class_ids = result.boxes.cls  # Class IDs
+            confs = result.boxes.conf  # Confidence scores
+            
+            for box, class_id, conf in zip(boxes, class_ids, confs):
+                if conf < 0.5:  # Skip low-confidence detections
+                    continue
+                
+                class_name = model.names[int(class_id)]  # Get class name
+                
+                if class_name in ["box", "suitcase", "handbag", "backpack"]:  # Consider as a package
+                    x1, y1, x2, y2 = map(int, box)
+                    
+                    # Check if detected package overlaps with a moving contour
+                    for cnt in large_contours:
+                        cx, cy, cw, ch = cv2.boundingRect(cnt)
+                        if x1 < cx < x2 and y1 < cy < y2:
+                            # Draw bounding box for detected package
+                            cv2.rectangle(frame_out, (x1, y1), (x2, y2), (0, 255, 0), 3)
+                            cv2.putText(frame_out, f"{class_name} ({conf:.2f})", (x1, y1 - 10),
+                                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
 
         frame_out_final = cv2.cvtColor(frame_out, cv2.COLOR_BGR2RGB)
         out.write(frame_out)
 
+        # Display every 12th frame
         if not count % 12:
             yield frame_out_final, None
         count += 1
@@ -82,9 +94,9 @@ def vid_inf(vid_path, contour_thresh):
     cv2.destroyAllWindows()
     yield None, output_video
 
-# Gradio interface
+# Gradio Interface
 input_video = gr.Video(label="Input Video")
-contour_thresh = gr.Slider(0, 10000, value=4, label="Contour Threshold", info="Adjust the threshold based on package size.")
+contour_thresh = gr.Slider(0, 10000, value=500, label="Contour Threshold")
 output_frames = gr.Image(label="Output Frames")
 output_video_file = gr.Video(label="Output Video")
 
@@ -92,11 +104,13 @@ app = gr.Interface(
     fn=vid_inf,
     inputs=[input_video, contour_thresh],
     outputs=[output_frames, output_video_file],
-    title="Package Tracking using YOLOv8 & Motion Detection",
-    description="A smart video analysis tool that uses YOLOv8 to track packages while ignoring human movement.",
+    title="Package Tracking using Motion Detection & Object Detection",
+    description="Detects and tracks moving packages while ignoring humans using YOLOv8 and OpenCV.",
     allow_flagging="never",
     examples=[["./sample/car.mp4", "1000"], ["./sample/motion_test.mp4", "5000"], ["./sample/home.mp4", "4500"]],
     cache_examples=False,
 )
+
 app.queue().launch()
 
+