Spaces:
Build error
Build error
| 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) | |