Spaces:
Sleeping
Sleeping
File size: 3,974 Bytes
7d1f745 dbc9309 970d1a0 eaf5124 970d1a0 eaf5124 4f7ba4c eaf5124 7d1f745 970d1a0 eaf5124 970d1a0 eaf5124 970d1a0 eaf5124 970d1a0 eaf5124 970d1a0 eaf5124 970d1a0 eaf5124 970d1a0 eaf5124 970d1a0 dbc9309 970d1a0 eaf5124 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import gradio as gr
import torch
import cv2
import os
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
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()
|