Spaces:
Build error
Build error
import gradio as gr | |
from ultralytics import YOLO | |
import cv2 | |
from deep_sort_realtime.deepsort_tracker import DeepSort | |
import tempfile | |
# Initialize YOLO model | |
model = YOLO("yolov8l.pt") # Load YOLOv8 model | |
tracker = DeepSort(max_age=30, n_init=3, nn_budget=100) | |
def count_people_in_video(video_file): | |
cap = cv2.VideoCapture(video_file) # Load video | |
total_ids = set() # Track unique IDs | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
# Run YOLO inference on the frame | |
results = model(frame) | |
detections = [] | |
# Parse YOLO detections | |
for result in results: | |
for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf): | |
if result.names[int(cls)] == "person" and conf > 0.5: # Detect "person" class | |
x1, y1, x2, y2 = map(int, box) | |
bbox = [x1, y1, x2 - x1, y2 - y1] # Convert to [x, y, width, height] | |
detections.append((bbox, conf, "person")) | |
# Update DeepSORT tracker with detections | |
tracks = tracker.update_tracks(detections, frame=frame) | |
# Add unique IDs from confirmed tracks | |
for track in tracks: | |
if track.is_confirmed(): | |
total_ids.add(track.track_id) | |
cap.release() | |
return len(total_ids) | |
# Gradio Interface | |
def process_video(video_file): | |
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file: | |
temp_file.write(video_file.read()) | |
temp_file.flush() | |
total_people = count_people_in_video(temp_file.name) | |
return f"Total unique people in the video: {total_people}" | |
interface = gr.Interface( | |
fn=process_video, | |
inputs=gr.Video(label="Upload a Video"), | |
outputs="text", | |
title="Person Counting with YOLOv8 and DeepSORT", | |
description="Upload a video to count the number of unique people using YOLOv8 and DeepSORT." | |
) | |
if __name__ == "__main__": | |
interface.launch() | |