Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import cv2 | |
| import os | |
| from ultralytics import YOLO | |
| import spaces | |
| # 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 | |
| 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}") | |
| 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() | |
| # Ensure output file exists before returning | |
| if not os.path.exists(output_full_path): | |
| raise FileNotFoundError(f"Output video not found: {output_full_path}") | |
| print(f"Processed video saved at: {output_full_path}") | |
| return output_full_path | |
| def process_video(video): | |
| try: | |
| print(f"Received video: {video}") | |
| detector = CrowdDetection() | |
| output_video = detector.detect_crowd(video) | |
| if not os.path.exists(output_video): # Ensure output file exists | |
| raise FileNotFoundError(f"Output video does not exist: {output_video}") | |
| 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 | |
| # 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" | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() | |