File size: 3,414 Bytes
d0e540e 875bfcc d0e540e 97bc949 d0e540e 0303123 97bc949 d0e540e 97bc949 4f0c296 d0e540e 4f0c296 d0e540e df95b60 d0e540e 4f0c296 |
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 |
import gradio as gr
from ultralytics import YOLO
import cv2
from deep_sort_realtime.deepsort_tracker import DeepSort
import logging
def initialize_tracker(max_age=30, n_init=3, nn_budget=100):
return DeepSort(max_age=max_age, n_init=n_init, nn_budget=nn_budget)
def detect_people(model, frame, confidence_threshold=0.5):
results = model(frame, device="cpu") # Force CPU if CUDA is unavailable
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 > confidence_threshold:
x1, y1, x2, y2 = map(int, box)
bbox = [x1, y1, x2 - x1, y2 - y1]
detections.append((bbox, conf, "person"))
return detections
def count_people_in_video(video_path, model_path="setosys_ppl_in_video_small_v1.pt", confidence_threshold=0.5):
logging.basicConfig(level=logging.INFO)
cap = cv2.VideoCapture(video_path)
model = YOLO(model_path) # Auto-detect device during model loading
tracker = initialize_tracker()
total_ids = set()
frame_count = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame_count += 1
logging.info(f"Processing frame {frame_count}")
detections = detect_people(model, frame, confidence_threshold)
tracks = tracker.update_tracks(detections, frame=frame)
for track in tracks:
if track.is_confirmed():
total_ids.add(track.track_id)
cap.release()
return len(total_ids)
# Initialize YOLO model
model = YOLO("setosys_ppl_in_video_small_v1.pt") # Load model
#tracker = DeepSort(max_age=30, n_init=3, nn_budget=100)
def count_people_in_video_old(video_path):
cap = cv2.VideoCapture(video_path) # 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):
# The `video_file` is a path to the temporary file
total_people = count_people_in_video(video_file)
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 in a Video",
description="Upload a video to count the number of unique people using YOLOv8 and DeepSORT."
)
if __name__ == "__main__":
interface.launch(share=True)
|