Update app.py

app.py

@@ -1,164 +1,199 @@
-import gradio as gr
+import time
+import torch
 import cv2
 import numpy as np
-from tensorflow.keras.models import load_model
-from sklearn.preprocessing import StandardScaler
+import gradio as gr
 from ultralytics import YOLO
+from deep_sort.utils.parser import get_config
+from deep_sort.deep_sort import DeepSort
 
-# Load models
-lstm_model = load_model('suspicious_activity_model.h5')
-yolo_model = YOLO('yolov8n-pose.pt')  # Ensure this model supports keypoint detection
-scaler = StandardScaler()
+# Initialize YOLO and DeepSort
+deep_sort_weights = 'ckpt.t7'
+tracker = DeepSort(model_path=deep_sort_weights, max_age=80)
+model = YOLO("person_gun.pt")
 
-def extract_keypoints(frame):
-    """
-    Extracts normalized keypoints from a frame using YOLO pose model.
-    """
-    results = yolo_model(frame, verbose=False)
-    for r in results:
-        if r.keypoints is not None and len(r.keypoints) > 0:
-            # Extract the first detected person's keypoints
-            keypoints = r.keypoints.xyn.tolist()[0]  # Use the first person's keypoints
-            flattened_keypoints = [kp for keypoint in keypoints for kp in keypoint[:2]]  # Flatten x, y values
-            return flattened_keypoints
-    return None  # Return None if no keypoints are detected
-
-def process_input(input_media):
-    """
-    Process either a video or an image for suspicious activity detection
-    """
-    # Determine if input is a video or image path
-    is_video = input_media.lower().endswith(('.mp4', '.avi', '.mov'))
-    
-    if is_video:
-        return process_video(input_media)
-    else:
-        return process_image(input_media)
+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 = []
+        
+        # Detection parameters
+        self.running_threshold = 0.5
+        self.fps = 30  # Default FPS
 
-def process_image(image_path):
-    """
-    Process a single image for suspicious activity detection
-    """
-    # Read the image
-    frame = cv2.imread(image_path)
-    
-    # Perform YOLO detection
-    results = yolo_model(frame, verbose=False)
-    for box in results[0].boxes:
-        cls = int(box.cls[0])  # Class ID
-        confidence = float(box.conf[0])
+    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 only (class_id 0 for 'person')
-        if cls == 0 and confidence > 0.5:
-            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates
-            
-            # Extract ROI for classification
-            roi = frame[y1:y2, x1:x2]
-            if roi.size > 0:
-                # Preprocess ROI to extract keypoints
-                keypoints = extract_keypoints(roi)
-                
-                if keypoints is not None and len(keypoints) > 0:
-                    # Standardize and reshape keypoints for LSTM input
-                    keypoints_scaled = scaler.fit_transform([keypoints])  # Standardize features
-                    keypoints_reshaped = keypoints_scaled.reshape((1, 1, len(keypoints)))  # Reshape for LSTM
-                    
-                    # Predict with LSTM model
-                    prediction = (lstm_model.predict(keypoints_reshaped) > 0.5).astype(int)[0][0]
-                    
-                    # Draw bounding box and label
-                    color = (0, 0, 255) if prediction == 1 else (0, 255, 0)
-                    label = 'Suspicious' if prediction == 1 else 'Normal'
-                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-                else:
-                    print("No valid keypoints detected for ROI. Skipping.")
+        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:
-                print("ROI size is zero. Skipping.")
-    
-    # Save the processed image
-    output_path = 'output_image.jpg'
-    cv2.imwrite(output_path, frame)
-    return output_path
+                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)
+            
+            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)
 
-def process_video(input_video):
+        # Release resources
+        cap.release()
+        out.release()
+
+        return 'output_detection.mp4', frame_info_list
+
+# Create Gradio interface
+detector = ObjectDetector()
+
+def detect_interface(input_media):
     """
-    Process video for suspicious activity detection
+    Gradio interface function for detection
     """
-    # Open video capture
-    cap = cv2.VideoCapture(input_video)
-    
-    # Prepare to save 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)
+    output_video, frame_info_list = detector.process_input(input_media)
     
-    # Create VideoWriter object
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter('output_video.mp4', fourcc, fps, (width, height))
+    # 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"
     
-    # Process each frame
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        # Perform YOLO detection
-        results = yolo_model(frame, verbose=False)
-        for box in results[0].boxes:
-            cls = int(box.cls[0])  # Class ID
-            confidence = float(box.conf[0])
-            
-            # Detect persons only (class_id 0 for 'person')
-            if cls == 0 and confidence > 0.5:
-                x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates
-                
-                # Extract ROI for classification
-                roi = frame[y1:y2, x1:x2]
-                if roi.size > 0:
-                    # Preprocess ROI to extract keypoints
-                    keypoints = extract_keypoints(roi)
-                    
-                    if keypoints is not None and len(keypoints) > 0:
-                        # Standardize and reshape keypoints for LSTM input
-                        keypoints_scaled = scaler.fit_transform([keypoints])  # Standardize features
-                        keypoints_reshaped = keypoints_scaled.reshape((1, 1, len(keypoints)))  # Reshape for LSTM
-                        
-                        # Predict with LSTM model
-                        prediction = (lstm_model.predict(keypoints_reshaped) > 0.5).astype(int)[0][0]
-                        
-                        # Draw bounding box and label
-                        color = (0, 0, 255) if prediction == 1 else (0, 255, 0)
-                        label = 'Suspicious' if prediction == 1 else 'Normal'
-                        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-                        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-                    else:
-                        print("No valid keypoints detected for ROI. Skipping frame.")
-                else:
-                    print("ROI size is zero. Skipping frame.")
-        
-        # Write processed frame to output video
-        out.write(frame)
-    
-    # Release resources
-    cap.release()
-    out.release()
-    
-    return 'output_video.mp4'
+    return output_video, summary
 
-# Create Gradio interface
+# Gradio Interface
 iface = gr.Interface(
-    fn=process_input,
+    fn=detect_interface,
     inputs=[
-        gr.File(label="Upload Image or Video", 
-                file_types=['image', 'video'], 
-                type="filepath")
+        gr.File(label="Upload Video", type="filepath"),
+        gr.Webcam(label="Or Use Webcam")
     ],
     outputs=[
-        gr.File(label="Processed Media")
+        gr.Video(label="Processed Video"),
+        gr.Textbox(label="Detection Summary")
     ],
-    title="Suspicious Activity Detection",
-    description="Upload an image or video to detect suspicious activities using YOLO and LSTM models. Suspicious activities will be marked with red bounding boxes, normal activities with green."
+    title="Object Detection with Tracking",
+    description="Upload a video or use webcam for real-time object detection and tracking"
 )
 
 # Launch the interface