Update app.py

app.py

@@ -1,107 +1,93 @@
 import gradio as gr
-import torch
 import cv2
+import numpy as np
 import os
+import torch
 from ultralytics import YOLO
-import spaces  
-
-@spaces.GPU  # Ensures GPU is allocated for this function
-
-class CrowdDetection:
-    def __init__(self, yolo_model_path="yolov8n.pt", crowd_threshold=10):
-        # Determine the best available device
-        if torch.cuda.is_available():
-            self.device = torch.device("cuda")
-            print(f"Using CUDA Device: {torch.cuda.get_device_name(0)}")
-        else:
-            self.device = torch.device("cpu")
-            print("Using CPU as no CUDA device is available")
-        
-        try:
-            # Load the YOLO model on the selected device
-            self.model = YOLO(yolo_model_path).to(self.device)
-        except Exception as e:
-            print(f"Error loading YOLO model: {e}")
-            self.model = None
-        
-        self.crowd_threshold = crowd_threshold
+import spaces  # Import ZeroGPU for Hugging Face Spaces
 
-    def detect_crowd(self, video_path):
-        if self.model is None:
-            raise ValueError("YOLO model failed to load")
-        
-        cap = cv2.VideoCapture(video_path)
-        
-        # Ensure video is opened successfully
-        if not cap.isOpened():
-            raise ValueError(f"Failed to open video: {video_path}")
+@spaces.GPU  # Ensures GPU is allocated during execution
+def process_video(video_path):
+    """Process video using YOLOv8 for crowd detection."""
+    
+    # Check if CUDA is available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"🔍 Using device: {device}")
 
-        output_path = "output_crowd.mp4"
-        output_full_path = os.path.abspath(output_path)  # Convert to absolute path
-        
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        out = cv2.VideoWriter(output_full_path, fourcc, int(cap.get(cv2.CAP_PROP_FPS)),
-                              (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
-        
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            
-            # Perform detection
-            results = self.model(frame)
-            person_count = 0
-            
-            for result in results:
-                boxes = result.boxes.xyxy.cpu().numpy()
-                classes = result.boxes.cls.cpu().numpy()
-                
-                for box, cls in zip(boxes, classes):
-                    if int(cls) == 0:  # YOLO class ID 0 = "person"
-                        person_count += 1
-                        x1, y1, x2, y2 = map(int, box)
-                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                        cv2.putText(frame, "Person", (x1, y1 - 10),
-                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-            
-            alert_text = "Crowd Alert!" if person_count > self.crowd_threshold else f"People: {person_count}"
-            cv2.putText(frame, alert_text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
-                        (0, 0, 255) if person_count > self.crowd_threshold else (0, 255, 0), 2)
-            
-            out.write(frame)
-        
-        cap.release()
-        out.release()
+    # Load YOLOv8 model on GPU
+    model = YOLO("yolov8n.pt").to(device)
+
+    # Read input video
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError(f"❌ Failed to open video: {video_path}")
 
-        # Ensure output file exists before returning
-        if not os.path.exists(output_full_path):
-            raise FileNotFoundError(f"Output video not found: {output_full_path}")
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-        print(f"Processed video saved at: {output_full_path}")
-        return output_full_path
+    print(f"🎥 Video details - FPS: {fps}, Width: {width}, Height: {height}")
 
-def process_video(video):
-    try:
-        print(f"Received video: {video}")
+    # Define output video path
+    output_path = "output_crowd.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    CROWD_THRESHOLD = 10  # Define crowd limit for alerts
+    frame_count = 0
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break  # End of video
         
-        detector = CrowdDetection()
-        output_video = detector.detect_crowd(video)
+        frame_count += 1
+
+        # Run YOLO inference on the frame
+        results = model(frame)
+
+        # Count detected persons
+        person_count = 0
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy()
+            classes = result.boxes.cls.cpu().numpy()
+
+            for box, cls in zip(boxes, classes):
+                if int(cls) == 0:  # YOLO class ID 0 = "person"
+                    person_count += 1
+                    x1, y1, x2, y2 = map(int, box)
+
+                    # Draw bounding box for persons
+                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                    cv2.putText(frame, "Person", (x1, y1 - 10),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+        # Display count on frame
+        alert_text = "Crowd Alert!" if person_count > CROWD_THRESHOLD else f"People: {person_count}"
+        cv2.putText(frame, alert_text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
+                    (0, 0, 255) if person_count > CROWD_THRESHOLD else (0, 255, 0), 2)
+
+        out.write(frame)  # Save frame to output video
+
+    cap.release()
+    out.release()
+
+    if frame_count == 0:
+        raise ValueError("❌ No frames were processed!")
 
-        if not os.path.exists(output_video):  # Ensure output file exists
-            raise FileNotFoundError(f"Output video does not exist: {output_video}")
+    if not os.path.exists(output_path):
+        raise FileNotFoundError(f"❌ Output video not found: {output_path}")
 
-        print(f"Returning processed video: {output_video}")
-        return output_video
-    except Exception as e:
-        print(f"Video processing error: {e}")
-        return None  # Prevent crashing the app
+    print(f"✅ Processed video saved at: {output_path}")
+    return output_path
 
-# Gradio Interface for Hugging Face Spaces
+# Gradio Interface
 interface = gr.Interface(
     fn=process_video,
     inputs=gr.Video(label="Upload Video"),
     outputs=gr.Video(label="Processed Video"),
-    title="Crowd Detection using YOLOv8"
+    title="Crowd Detection with YOLOv8",
+    description="Upload a video, and YOLOv8 will detect and count people. If the crowd exceeds 10 people, a warning will be displayed."
 )
 
 if __name__ == "__main__":