Update app.py

app.py

@@ -1,200 +1,128 @@
-import time
-import torch
 import cv2
+import torch
 import numpy as np
 import gradio as gr
 from ultralytics import YOLO
-from deep_sort.utils.parser import get_config
-from deep_sort.deep_sort import DeepSort
-
-# Initialize YOLO and DeepSort
-deep_sort_weights = 'ckpt.t7'
-tracker = DeepSort(model_path=deep_sort_weights, max_age=80)
-model = YOLO("person_gun.pt")
+from deep_sort_realtime.deep_sort import DeepSort
 
-class ObjectDetector:
-    def __init__(self):
-        # Tracking variables
-        self.unique_track_ids = set()
-        self.track_labels = {}
-        self.track_times = {}
-        self.track_positions = {}
-        self.running_counters = {}
-        self.alert_person_ids = []
+class ObjectTracker:
+    def __init__(self, person_model_path='yolov8n.pt'):
+        """
+        Initialize object tracker with YOLO and DeepSort
+        """
+        # Load YOLO model for person detection
+        self.model = YOLO(person_model_path)
+        
+        # Initialize DeepSort tracker
+        self.tracker = DeepSort(
+            max_age=30,  # Tracks can be lost for up to 30 frames
+            n_init=3,    # Number of consecutive detections before track is confirmed
+        )
         
-        # Detection parameters
-        self.running_threshold = 0.5
-        self.fps = 30  # Default FPS
+        # Tracking statistics
+        self.person_count = 0
+        self.tracking_data = {}
 
     def process_frame(self, frame):
         """
         Process a single frame for object detection and tracking
         """
-        # Reset alert tracking for this frame
-        self.alert_person_ids.clear()
-        og_frame = frame.copy()
-
-        # Detect persons
-        results = model(frame, device=0, classes=0, conf=0.75)
+        # Detect persons using YOLO
+        results = self.model(frame, classes=[0], conf=0.5)
         
-        for result in results:
-            boxes = result.boxes
-            cls = boxes.cls.tolist()
-            conf = boxes.conf
-            xywh = boxes.xywh
-
-        pred_cls = np.array(cls)
-        conf = conf.detach().cpu().numpy()
-        bboxes_xywh = xywh.cpu().numpy()
-
-        # Update tracking
-        tracks = tracker.update(bboxes_xywh, conf, og_frame)
-        active_track_ids = set()
-
-        # Reset running status
-        new_running_status = "No Running"
-
-        for track in tracker.tracker.tracks:
-            track_id = track.track_id
-            x1, y1, x2, y2 = track.to_tlbr()
-            w = x2 - x1
-            h = y2 - y1
-
-            # Define color for bounding box
-            red_color = (0, 0, 255)
-            blue_color = (255, 0, 0)
-            green_color = (0, 255, 0)
-            color_id = track_id % 3
-            color = red_color if color_id == 0 else blue_color if color_id == 1 else green_color
-            cv2.rectangle(og_frame, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
-
-            # Initialize tracking for new tracks
-            if track_id not in self.track_labels:
-                self.track_labels[track_id] = "Person"
-                self.track_times[track_id] = 0
-                self.track_positions[track_id] = (x1, y1)
-                self.running_counters[track_id] = 0
-
-            self.track_times[track_id] += 1
-            prev_x1, prev_y1 = self.track_positions[track_id]
-            displacement = np.sqrt((x1 - prev_x1) ** 2 + (y1 - prev_y1) ** 2)
-
-            # Calculate speed
-            speed = displacement / self.fps if self.fps > 0 else 0
-
-            self.track_positions[track_id] = (x1, y1)
-
-            # Detect running
-            if speed > self.running_threshold and w * h > 5000:
-                self.running_counters[track_id] += 1
-                if self.running_counters[track_id] > self.fps/2:
-                    self.track_labels[track_id] = "Running"
-                    new_running_status = "Running Detected"
-            else:
-                self.running_counters[track_id] = 0
-                self.track_labels[track_id] = "Person"
-
-            # Track time and potential alerts
-            total_seconds = self.track_times[track_id] / self.fps if self.fps > 0 else 0
-            minutes = int(total_seconds // 60)
-            seconds = int(total_seconds % 60)
-
-            # Trigger alert for prolonged stay
-            if total_seconds > 60 and track_id not in self.alert_person_ids:
-                self.alert_person_ids.append(track_id)
-
-            # Add label to frame
-            cv2.putText(og_frame, f"{self.track_labels[track_id]} {minutes:02}:{seconds:02}", 
-                       (int(x1) + 10, int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
+        # Extract bounding boxes and confidences
+        detections = []
+        for r in results:
+            boxes = r.boxes
+            for box in boxes:
+                # Convert to [x, y, w, h] format for DeepSort
+                x1, y1, x2, y2 = box.xyxy[0]
+                bbox = [x1.item(), y1.item(), (x2-x1).item(), (y2-y1).item()]
+                conf = box.conf.item()
+                detections.append((bbox, conf))
+        
+        # Update tracks
+        if detections:
+            tracks = self.tracker.update_tracks(
+                detections, 
+                frame=frame
+            )
             
-            active_track_ids.add(track_id)
-
-        # Update unique track IDs
-        self.unique_track_ids.intersection_update(active_track_ids)
-        self.unique_track_ids.update(active_track_ids)
-
-        # Prepare result dictionary
-        result_info = {
-            'person_count': len(self.unique_track_ids),
-            'running_status': new_running_status,
-            'prolonged_stay_ids': list(self.alert_person_ids)
-        }
-
-        return og_frame, result_info
-
-    def process_input(self, input_media):
-        """
-        Process either video or webcam input
-        """
-        # Determine input type
-        if isinstance(input_media, str):  # Video file path
-            cap = cv2.VideoCapture(input_media)
-        else:  # Webcam input
-            cap = cv2.VideoCapture(0)
-
-        # Prepare output video
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        fps = cap.get(cv2.CAP_PROP_FPS) or 30
-        self.fps = fps
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter('output_detection.mp4', fourcc, fps, (width, height))
-
-        # Processing loop
-        frame_info_list = []
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-
-            # Process frame
-            processed_frame, frame_info = self.process_frame(frame)
-            out.write(processed_frame)
-            frame_info_list.append(frame_info)
-
-        # Release resources
-        cap.release()
-        out.release()
-
-        return 'output_detection.mp4', frame_info_list
-
-# Create Gradio interface
-detector = ObjectDetector()
+            # Annotate frame with tracking information
+            for track in tracks:
+                if not track.is_confirmed():
+                    continue
+                
+                track_id = track.track_id
+                ltrb = track.to_ltrb()
+                
+                # Draw bounding box
+                cv2.rectangle(
+                    frame, 
+                    (int(ltrb[0]), int(ltrb[1])), 
+                    (int(ltrb[2]), int(ltrb[3])), 
+                    (0, 255, 0), 
+                    2
+                )
+                
+                # Add track ID
+                cv2.putText(
+                    frame, 
+                    f'ID: {track_id}', 
+                    (int(ltrb[0]), int(ltrb[1]-10)), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 
+                    0.9, 
+                    (0, 255, 0), 
+                    2
+                )
+        
+        return frame
 
-def detect_interface(input_media):
+def process_video(input_video):
     """
-    Gradio interface function for detection
+    Main video processing function for Gradio
     """
-    output_video, frame_info_list = detector.process_input(input_media)
+    # Initialize tracker
+    tracker = ObjectTracker()
+    
+    # Open input video
+    cap = cv2.VideoCapture(input_video)
+    
+    # Prepare output video writer
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
     
-    # Generate text summary
-    summary = "Detection Summary:\n"
-    if frame_info_list:
-        # Take the last frame's information
-        last_frame_info = frame_info_list[-1]
-        summary += f"Total Persons Detected: {last_frame_info['person_count']}\n"
-        summary += f"Running Status: {last_frame_info['running_status']}\n"
-        if last_frame_info['prolonged_stay_ids']:
-            summary += f"Prolonged Stay Detected - Person IDs: {last_frame_info['prolonged_stay_ids']}"
-        else:
-            summary += "No Prolonged Stay Detected"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter('output_tracked.mp4', fourcc, fps, (width, height))
     
-    return output_video, summary
+    # Process video frames
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        # Process and annotate frame
+        processed_frame = tracker.process_frame(frame)
+        
+        # Write processed frame
+        out.write(processed_frame)
+    
+    # Release resources
+    cap.release()
+    out.release()
+    
+    return 'output_tracked.mp4'
 
-# Gradio Interface
+# Create Gradio interface
 iface = gr.Interface(
-    fn=detect_interface,
-    inputs=[
-        gr.File(label="Upload Video", type="filepath"),
-        gr.Webcam(label="Or Use Webcam")
-    ],
-    outputs=[
-        gr.Video(label="Processed Video"),
-        gr.Textbox(label="Detection Summary")
-    ],
-    title="Object Detection with Tracking",
-    description="Upload a video or use webcam for real-time object detection and tracking"
+    fn=process_video,
+    inputs=gr.Video(label="Upload Video for Tracking"),
+    outputs=gr.Video(label="Tracked Video"),
+    title="Person Tracking with YOLO and DeepSort",
+    description="Upload a video to track and annotate person movements"
 )
 
 # Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()