import gradio as gr import cv2 import json import tempfile import os from ultralytics import YOLO import numpy as np from collections import defaultdict from typing import Dict, List, Tuple, Any class HumanTracker: def __init__(self): # Load YOLOv11 model - using the nano version for faster processing # You can change to yolo11s.pt, yolo11m.pt, yolo11l.pt, or yolo11x.pt for better accuracy self.model = YOLO("yolo11n.pt") def calculate_center(self, x1: float, y1: float, x2: float, y2: float) -> Tuple[float, float]: """Calculate center coordinates from bounding box coordinates.""" center_x = (x1 + x2) / 2 center_y = (y1 + y2) / 2 return center_x, center_y def process_video(self, video_path: str, progress_callback=None) -> Dict[str, Any]: """ Process video file and extract human tracking data. Args: video_path: Path to the input video file progress_callback: Optional callback function for progress updates Returns: Dictionary containing processed tracking data in the required JSON format """ cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f"Could not open video file: {video_path}") total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = cap.get(cv2.CAP_PROP_FPS) frame_data = {} id_mapping = {} # Maps original YOLO IDs to simplified sequential IDs next_person_id = 1 print(f"Processing video: {total_frames} frames at {fps} FPS") # Process video with YOLO tracking # Using stream=True for memory efficiency with large videos results = self.model.track( video_path, classes=[0], # Only detect humans (class 0) persist=True, # Enable tracking stream=True, verbose=False ) frame_count = 0 for result in results: if progress_callback: progress = (frame_count + 1) / total_frames progress_callback(progress, f"Processing frame {frame_count + 1}/{total_frames}") # Check if any detections exist if result.boxes is not None and len(result.boxes) > 0: # Extract bounding boxes, track IDs, and confidences boxes = result.boxes.xyxy.cpu().numpy() # x1, y1, x2, y2 format track_ids = result.boxes.id confidences = result.boxes.conf.cpu().numpy() if track_ids is not None: track_ids = track_ids.int().cpu().numpy() people_in_frame = [] for box, track_id, confidence in zip(boxes, track_ids, confidences): x1, y1, x2, y2 = box # Map original YOLO ID to simplified sequential ID if track_id not in id_mapping: id_mapping[track_id] = next_person_id next_person_id += 1 person_id = id_mapping[track_id] # Calculate center coordinates center_x, center_y = self.calculate_center(x1, y1, x2, y2) # Create person data person_data = { "person_id": person_id, "center_x": float(center_x), "center_y": float(center_y), "confidence": float(confidence), "bbox": { "x1": float(x1), "y1": float(y1), "x2": float(x2), "y2": float(y2) } } people_in_frame.append(person_data) if people_in_frame: # Sort people by person_id for consistency people_in_frame.sort(key=lambda x: x["person_id"]) frame_data[frame_count] = people_in_frame frame_count += 1 cap.release() # Convert to the required JSON format frames_list = [] sorted_frames = sorted(frame_data.keys()) for frame_num in sorted_frames: frames_list.append({ "frame": frame_num, "people": frame_data[frame_num] }) # Create the final output structure output = { "metadata": { "total_frames": len(frames_list), "total_people": len(id_mapping), "video_info": { "fps": float(fps), "total_video_frames": total_frames }, "id_mapping": {str(original_id): simplified_id for original_id, simplified_id in id_mapping.items()} }, "frames": frames_list } return output def process_video_gradio(video_file, progress=gr.Progress()): """ Gradio interface function for processing videos. Args: video_file: Uploaded video file from Gradio progress: Gradio progress tracker Returns: Tuple of (JSON file path, status message, preview of results) """ if video_file is None: return None, "❌ Please upload a video file", "No video uploaded" try: # Initialize the tracker tracker = HumanTracker() # Create progress callback def update_progress(prog, msg): progress(prog, desc=msg) # Process the video progress(0.1, desc="Starting video processing...") results = tracker.process_video(video_file, update_progress) progress(0.9, desc="Generating JSON output...") # Create temporary JSON file with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: json.dump(results, f, indent=2) json_path = f.name # Create a preview of the results metadata = results["metadata"] total_frames = metadata["total_frames"] total_people = metadata["total_people"] preview = f""" 📊 **Processing Results:** - **Total frames with detections:** {total_frames} - **Unique people detected:** {total_people} - **Original video frames:** {metadata.get('video_info', {}).get('total_video_frames', 'N/A')} - **Video FPS:** {metadata.get('video_info', {}).get('fps', 'N/A'):.2f} 🆔 **ID Mapping:** {json.dumps(metadata["id_mapping"], indent=2)} 📋 **Sample Frame Data (first frame):** {json.dumps(results["frames"][:1] if results["frames"] else [], indent=2)} """ progress(1.0, desc="✅ Processing complete!") return ( json_path, f"✅ Successfully processed video! Detected {total_people} unique people across {total_frames} frames.", preview ) except Exception as e: error_msg = f"❌ Error processing video: {str(e)}" print(error_msg) return None, error_msg, f"Error details: {str(e)}" # Create the Gradio interface def create_interface(): with gr.Blocks( title="Dynamic Veme Processor", theme=gr.themes.Soft() ) as demo: gr.Markdown(""" # 🎯 Dynamic Veme Processor Upload a video to detect and track humans using YOLOv11. The app will: - 🔍 Detect humans in each frame - 🎯 Track individuals across frames with unique IDs - 📐 Extract bounding box coordinates and center points - 📁 Generate JSON output for text overlay positioning **Supported formats:** MP4, AVI, MOV, WEBM """) with gr.Row(): with gr.Column(scale=1): video_input = gr.Video( label="📹 Upload Video", height=400 ) process_btn = gr.Button( "🚀 Process Video", variant="primary", size="lg" ) with gr.Column(scale=1): json_output = gr.File( label="📁 Download JSON Results", file_count="single" ) status_output = gr.Textbox( label="📊 Status", value="Ready to process video...", interactive=False ) with gr.Row(): preview_output = gr.Textbox( label="👁️ Results Preview", lines=15, interactive=False, placeholder="Results preview will appear here after processing..." ) # Event handlers process_btn.click( fn=process_video_gradio, inputs=[video_input], outputs=[json_output, status_output, preview_output], show_progress=True ) # Example section gr.Markdown(""" ## 📋 Output Format The generated JSON file contains: - **metadata**: Video info, total people count, ID mappings - **frames**: Array of frame data with person detections Each person detection includes: - `person_id`: Unique identifier for tracking - `center_x`, `center_y`: Center coordinates for text overlay positioning - `confidence`: Detection confidence score - `bbox`: Full bounding box coordinates (x1, y1, x2, y2) """) return demo if __name__ == "__main__": # Create and launch the interface demo = create_interface() demo.launch( server_name="0.0.0.0", # Allow external access server_port=7860, share=False, # Set to True if you want a public link show_error=True )