Update app.py

app.py

@@ -1,20 +1,88 @@
 import gradio as gr
 import torch
-import spaces  # Import ZeroGPU helper
+import cv2
+from ultralytics import YOLO
+import spaces 
+@spaces.GPU  
 
-# Function to check GPU availability
-@spaces.GPU  # This tells Hugging Face to allocate a GPU when this function is called
-def check_gpu():
-    if torch.cuda.is_available():
-        return "✅ GPU is available! CUDA device: " + torch.cuda.get_device_name(0)
-    else:
-        return "❌ No GPU detected!"
 
-# Gradio UI
-iface = gr.Interface(
-    fn=check_gpu,  # Call the GPU-checking function
-    inputs=[],
-    outputs="text",
+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
+
+    def detect_crowd(self, video_path):
+        if self.model is None:
+            raise ValueError("YOLO model failed to load")
+        
+        cap = cv2.VideoCapture(video_path)
+        output_path = "output_crowd.mp4"
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        out = cv2.VideoWriter(output_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()
+        return output_path
+
+def process_video(video):
+    try:
+        detector = CrowdDetection()
+        output_video = detector.detect_crowd(video)
+        return output_video
+    except Exception as e:
+        print(f"Video processing error: {e}")
+        return None
+
+# Gradio Interface for Hugging Face Spaces
+interface = gr.Interface(
+    fn=process_video,
+    inputs=gr.Video(label="Upload Video"),
+    outputs=gr.Video(label="Processed Video"),
+    title="Crowd Detection using YOLOv8"
 )
 
-iface.launch()
+# Remove share=True for Hugging Face Spaces
+if __name__ == "__main__":
+    interface.launch()