Spaces:
Sleeping
Sleeping
File size: 2,361 Bytes
4f7ba4c |
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 |
import gradio as gr
import torch
import cv2
from ultralytics import YOLO
# Crowd Detection Class
class CrowdDetection:
def __init__(self, yolo_model_path="yolov8n.pt", crowd_threshold=10):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = YOLO(yolo_model_path).to(self.device) # Move to GPU
self.model.nms = 0.5 # Set NMS threshold
self.crowd_threshold = crowd_threshold
def detect_crowd(self, video_path):
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
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
# Gradio Function
def process_video(video):
detector = CrowdDetection()
return detector.detect_crowd(video)
# Gradio Interface
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()
|