Update app.py

app.py

@@ -5,88 +5,62 @@ from ultralytics import YOLO
 import cv2
 import gradio as gr
 import numpy as np
-import torch
 from ultralytics import YOLO
+from PIL import Image
 
-# Load YOLOv8 model (pre-trained)
-model = YOLO("yolov8n.pt")  # Use a small YOLOv8 model for efficiency
+# Load YOLO model
+yolo_model = YOLO('yolov8n.pt')
+TARGET_LABEL = 'parcel'  # Define the target object label
 
-# Video Inference Function
 def vid_inf(vid_path, contour_thresh):
     cap = cv2.VideoCapture(vid_path)
-    
-    # Get video properties
-    frame_width = int(cap.get(3))
-    frame_height = int(cap.get(4))
+    frame_width, frame_height = int(cap.get(3)), 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 None, None
-
     count = 0
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        
-        # 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)
         large_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > contour_thresh]
 
+        # Run YOLO detection
+        results = yolo_model(frame)
         frame_out = frame.copy()
 
-        # 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
+            for id, box in enumerate(result.boxes.xyxy):
+                class_id = int(result.boxes.cls[id])
+                label = yolo_model.names[class_id]
+                conf = result.boxes.conf[id]
                 
-                if class_name in ["box", "suitcase", "handbag", "backpack"]:  # Consider as a package
+                if label == TARGET_LABEL and conf >= 0.5:
                     x1, y1, x2, y2 = map(int, box)
-                    
-                    # Check if detected package overlaps with a moving contour
+                    center = ((x1 + x2) // 2, (y1 + y2) // 2)
                     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)
-
+                        if cv2.contourArea(cnt) > 500:
+                            x, y, w, h = cv2.boundingRect(cnt)
+                            if x1 < x < x2 and y1 < y < y2:
+                                cv2.rectangle(frame_out, (x1, y1), (x2, y2), (0, 255, 255), 2)
+                                cv2.putText(frame_out, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 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
+            yield Image.fromarray(frame_out_final), None
         count += 1
 
     cap.release()
@@ -94,7 +68,7 @@ def vid_inf(vid_path, contour_thresh):
     cv2.destroyAllWindows()
     yield None, output_video
 
-# Gradio Interface
+# Gradio UI Setup
 input_video = gr.Video(label="Input Video")
 contour_thresh = gr.Slider(0, 10000, value=500, label="Contour Threshold")
 output_frames = gr.Image(label="Output Frames")
@@ -104,10 +78,9 @@ app = gr.Interface(
     fn=vid_inf,
     inputs=[input_video, contour_thresh],
     outputs=[output_frames, output_video_file],
-    title="Package Tracking using Motion Detection & Object Detection",
-    description="Detects and tracks moving packages while ignoring humans using YOLOv8 and OpenCV.",
+    title="YOLO Motion Detection - Parcel Focus",
+    description='A video analysis tool using YOLOv8 for parcel detection with motion tracking.',
     allow_flagging="never",
-    examples=[["./sample/car.mp4", "1000"], ["./sample/motion_test.mp4", "5000"], ["./sample/home.mp4", "4500"]],
     cache_examples=False,
 )