Update app.py

app.py

@@ -12,8 +12,16 @@ class CrowdDetection:
     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")
-        self.model = YOLO(model_path).to(self.device)  # Load model once
+        if not os.path.exists(model_path):
+            # Download the model if not present
+            from ultralytics import YOLO
+            self.model = YOLO("yolov8n.pt")  # This downloads the model automatically
+            self.model.save(model_path)  # Save locally
+        else:
+            self.model = YOLO(model_path)
+        self.model.to(self.device)
 
+    @spaces.GPU
     def detect_crowd(self, video_path):
         """Process video using YOLOv8 for crowd detection."""
         cap = cv2.VideoCapture(video_path)
@@ -28,6 +36,10 @@ class CrowdDetection:
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
 
+        if not out.isOpened():
+            cap.release()
+            raise ValueError(f"❌ Failed to initialize video writer for {output_path}")
+
         CROWD_THRESHOLD = 10
         frame_count = 0
 
@@ -75,8 +87,37 @@ class CrowdDetection:
         if not os.path.exists(output_path):
             raise FileNotFoundError(f"❌ Output video not found: {output_path}")
 
-        return output_path  # Return file path instead of video object
+        return output_path
+
+# Define Gradio interface function
+def process_video(video):
+    try:
+        detector = CrowdDetection()  # Instantiate inside to avoid pickling
+        output_path = detector.detect_crowd(video)
+        return "Crowd detection complete!", output_path
+    except Exception as e:
+        return f"Error: {str(e)}", None
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Crowd Detection with YOLOv8")
+    gr.Markdown("Upload a video to detect people and get crowd alerts (threshold: 10 people)")
+    
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Upload Video")
+            submit_btn = gr.Button("Detect Crowd")
+        with gr.Column():
+            status_output = gr.Textbox(label="Status")
+            video_output = gr.Video(label="Result")
+    
+    submit_btn.click(
+        fn=process_video,
+        inputs=[video_input],
+        outputs=[status_output, video_output]
+    )
 
+demo.launch(debug=True)
 
 
 class PeopleTracking: