hb-setosys's picture
Create app.py
d0e540e verified
raw
history blame
2.02 kB
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()