Update app.py

app.py

@@ -8,14 +8,19 @@ import spaces
 
 @spaces.GPU
 
+import cv2
+import torch
+import os
+from ultralytics import YOLO
+
 class CrowdDetection:
-    def __init__(self, model_path="yolov8n.pt", crowd_threshold=10):
+    def __init__(self, model_path="yolov8n.pt"):
+        """Initialize the YOLO model once to avoid PicklingError."""
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"🔍 Using device: {self.device}")
-        self.model = YOLO(model_path).to(self.device)
-        self.crowd_threshold = crowd_threshold
+        self.model = YOLO(model_path).to(self.device)  # Load model once
 
     def detect_crowd(self, video_path):
+        """Process video using YOLOv8 for crowd detection."""
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise ValueError(f"❌ Failed to open video: {video_path}")
@@ -23,38 +28,47 @@ class CrowdDetection:
         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"🎥 Video details - FPS: {fps}, Width: {width}, Height: {height}")
 
         output_path = "output_crowd.mp4"
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
 
+        CROWD_THRESHOLD = 10
         frame_count = 0
 
         while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
-                break
+                break  # End of video
             
             frame_count += 1
+
+            # Run YOLO inference on the frame
             results = self.model(frame)
-            person_count = 0
 
+            # Count detected persons
+            person_count = sum(
+                1 for result in results
+                for cls in result.boxes.cls.cpu().numpy() if int(cls) == 0
+            )
+
+            # Draw bounding boxes
             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:  # Class ID 0 = "person"
-                        person_count += 1
+                    if int(cls) == 0:  # Person class
                         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}"
+            # 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 > self.crowd_threshold else (0, 255, 0), 2)
+                        (0, 0, 255) if person_count > CROWD_THRESHOLD else (0, 255, 0), 2)
+
             out.write(frame)
 
         cap.release()
@@ -66,8 +80,9 @@ class CrowdDetection:
         if not os.path.exists(output_path):
             raise FileNotFoundError(f"❌ Output video not found: {output_path}")
 
-        print(f"✅ Processed video saved at: {output_path}")
-        return output_path
+        return output_path  # Return file path instead of video object
+
+
 
 class PeopleTracking:
     def __init__(self, yolo_model_path="yolov8n.pt"):