Spaces:
Running
Running
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from collections import defaultdict | |
| import tempfile | |
| import os | |
| class PersonCounter: | |
| def __init__(self, line_position=0.5): | |
| self.model = YOLO("yolov8n.pt") | |
| self.tracker = defaultdict(list) | |
| self.crossed_ids = set() | |
| self.line_position = line_position | |
| self.count = 0 | |
| def process_frame(self, frame): | |
| height, width = frame.shape[:2] | |
| line_y = int(height * self.line_position) | |
| # Draw counting line | |
| cv2.line(frame, (0, line_y), (width, line_y), (0, 255, 0), 2) | |
| # Run detection and tracking | |
| results = self.model.track(frame, persist=True, classes=[0]) | |
| if results[0].boxes.id is not None: | |
| boxes = results[0].boxes.xyxy.cpu().numpy() | |
| track_ids = results[0].boxes.id.cpu().numpy().astype(int) | |
| for box, track_id in zip(boxes, track_ids): | |
| # Draw bounding box | |
| cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), | |
| (255, 0, 0), 2) | |
| # Get feet position | |
| center_x = (box[0] + box[2]) / 2 | |
| feet_y = box[3] | |
| # Draw tracking point | |
| cv2.circle(frame, (int(center_x), int(feet_y)), 5, (0, 255, 255), -1) | |
| # Store tracking history | |
| if track_id in self.tracker: | |
| prev_y = self.tracker[track_id][-1] | |
| # Check if person has crossed the line | |
| if prev_y < line_y and feet_y >= line_y and track_id not in self.crossed_ids: | |
| self.crossed_ids.add(track_id) | |
| self.count += 1 | |
| # Draw crossing indicator | |
| cv2.circle(frame, (int(center_x), int(line_y)), 8, (0, 0, 255), -1) | |
| self.tracker[track_id] = [feet_y] | |
| # Draw count with background | |
| count_text = f"Count: {self.count}" | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| font_scale = 1.5 | |
| thickness = 3 | |
| (text_width, text_height), _ = cv2.getTextSize(count_text, font, font_scale, thickness) | |
| cv2.rectangle(frame, (10, 10), (20 + text_width, 20 + text_height), | |
| (0, 0, 0), -1) | |
| cv2.putText(frame, count_text, (15, 15 + text_height), | |
| font, font_scale, (0, 255, 0), thickness) | |
| return frame | |
| def process_video(video_path, progress=gr.Progress()): | |
| # Create temp directory for output | |
| temp_dir = tempfile.mkdtemp() | |
| output_path = os.path.join(temp_dir, "result.mp4") | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| raise ValueError("Could not open video file") | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height)) | |
| counter = PersonCounter(line_position=0.5) | |
| for frame_idx in progress.tqdm(range(total_frames)): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| processed_frame = counter.process_frame(frame) | |
| writer.write(processed_frame) | |
| cap.release() | |
| writer.release() | |
| return output_path, f"Final count: {counter.count} people entered" | |
| # Create Gradio interface | |
| demo = gr.Interface( | |
| fn=process_video, | |
| inputs=gr.Video(label="Upload a video file"), | |
| outputs=[ | |
| gr.Video(label="Processed Video"), | |
| gr.Textbox(label="Results") | |
| ], | |
| title="Store Entry People Counter", | |
| description="Upload a video to count the number of people entering through a line. The green line represents the counting threshold, blue boxes show detected people, and the counter increases when someone crosses the line from top to bottom.", | |
| examples=[], | |
| cache_examples=False | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |