import warnings warnings.filterwarnings("ignore") import gradio as gr import cv2 import numpy as np import json import os from datetime import datetime from ultralytics import YOLO from insightface.app import FaceAnalysis import torchreid import torch import logging import shutil import tempfile import uuid # ========== Logging Configuration ========== logging.basicConfig( level=logging.INFO, format='[%(asctime)s] [%(levelname)s] %(message)s', handlers=[ logging.FileHandler("app.log"), logging.StreamHandler() ] ) logger = logging.getLogger(__name__) # ========== Configuration ========== DETECTION_THRESHOLD = 0.75 # Create output directory for Gradio OUTPUT_DIR = os.path.join(os.getcwd(), "outputs") os.makedirs(OUTPUT_DIR, exist_ok=True) # ========== Video Processing Class ========== class VideoProcessor: def __init__(self): try: self.model = YOLO('detection.pt') self.face_app = FaceAnalysis(name='buffalo_l', providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) self.face_app.prepare(ctx_id=0) self.reid_extractor = torchreid.utils.FeatureExtractor( model_name='osnet_x0_25', model_path=None, device='cuda' if torch.cuda.is_available() else 'cpu' ) self.models_loaded = True logger.info("Models loaded successfully.") except Exception as e: logger.exception("Model loading failed.") self.models_loaded = False self.reset_tracking() def reset_tracking(self): self.known_embeddings = [] self.known_ids = [] self.next_global_id = 1 self.track_to_global = {} self.tracking_data = { "metadata": { "total_frames": 0, "total_people": 0, "id_mapping": {} }, "frames": [] } logger.info("Tracking state reset.") def extract_embeddings(self, person_crop): face_embedding, body_embedding = None, None try: faces = self.face_app.get(person_crop) if faces: face_embedding = faces[0].embedding except Exception: logger.debug("Face embedding failed.") try: body_input = cv2.resize(person_crop, (128, 256)) body_input = cv2.cvtColor(body_input, cv2.COLOR_BGR2RGB) body_embedding = self.reid_extractor(body_input)[0].cpu().numpy() except Exception: logger.debug("Body embedding failed.") if face_embedding is not None and body_embedding is not None: return np.concatenate((face_embedding, body_embedding)).astype(np.float32) elif face_embedding is not None: return face_embedding.astype(np.float32) elif body_embedding is not None: return body_embedding.astype(np.float32) return None def assign_global_id(self, embedding, track_id): if embedding is None: return self.track_to_global.get(track_id, f"T{track_id}") match_found = False if self.known_embeddings: matching_embeddings = [ (emb, gid) for emb, gid in zip(self.known_embeddings, self.known_ids) if emb.shape[0] == embedding.shape[0] ] if matching_embeddings: embs, gids = zip(*matching_embeddings) embs = np.array(embs) sims = np.dot(embs, embedding) / ( np.linalg.norm(embs, axis=1) * np.linalg.norm(embedding) + 1e-6 ) best_match = np.argmax(sims) if sims[best_match] > 0.6: global_id = gids[best_match] match_found = True if not match_found: global_id = self.next_global_id self.next_global_id += 1 self.known_embeddings.append(embedding) self.known_ids.append(global_id) if track_id is not None: self.track_to_global[track_id] = global_id return global_id def process_video(self, input_video_path, progress_callback=None): if not self.models_loaded: raise Exception("Models not loaded properly") self.reset_tracking() # Create output files with timestamp timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") unique_id = str(uuid.uuid4())[:8] # Use the OUTPUT_DIR instead of temp directory output_video_path = os.path.join(OUTPUT_DIR, f"tracked_video_{timestamp}_{unique_id}.mp4") output_json_path = os.path.join(OUTPUT_DIR, f"tracking_data_{timestamp}_{unique_id}.json") cap = cv2.VideoCapture(input_video_path) if not cap.isOpened(): raise Exception("Could not open video file") width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Use H.264 codec for better compatibility and add proper video codec fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Changed from 'mp4v' to 'H264' out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height)) # Verify video writer is properly initialized if not out.isOpened(): logger.warning("H264 codec failed, trying XVID") fourcc = cv2.VideoWriter_fourcc(*'XVID') output_video_path = output_video_path.replace('.mp4', '.avi') out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height)) if not out.isOpened(): logger.warning("XVID codec failed, trying mp4v") fourcc = cv2.VideoWriter_fourcc(*'H264') output_video_path = output_video_path.replace('.avi', '.mp4') out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height)) frame_count = 0 while True: ret, frame = cap.read() if not ret: break frame_count += 1 if progress_callback: progress_callback(frame_count / total_frames, f"Processing frame {frame_count}/{total_frames}") frame_data = {"frame": frame_count, "people": []} try: results = self.model.track( frame, tracker="bytetrack.yaml", persist=True, verbose=False, conf=DETECTION_THRESHOLD ) for result in results: if result.boxes is not None: boxes = result.boxes.xyxy.cpu().numpy() confidences = result.boxes.conf.cpu().numpy() track_ids = result.boxes.id.int().cpu().tolist() if result.boxes.id is not None else [None] * len(boxes) for box, conf, track_id in zip(boxes, confidences, track_ids): x1, y1, x2, y2 = map(int, box) person_crop = frame[y1:y2, x1:x2] if person_crop.size > 0: embedding = self.extract_embeddings(person_crop) global_id = self.assign_global_id(embedding, track_id) frame_data["people"].append({ "person_id": global_id, "center_x": (x1 + x2) / 2, "center_y": (y1 + y2) / 2, "confidence": float(conf), "bbox": {"x1": float(x1), "y1": float(y1), "x2": float(x2), "y2": float(y2)} }) cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(frame, f"ID {global_id}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) except Exception as e: logger.exception(f"Error processing frame {frame_count}.") self.tracking_data["frames"].append(frame_data) out.write(frame) cap.release() out.release() # Verify the output file was created and has content if not os.path.exists(output_video_path) or os.path.getsize(output_video_path) == 0: raise Exception("Output video file was not created properly") self.tracking_data["metadata"]["total_frames"] = frame_count self.tracking_data["metadata"]["total_people"] = len(set(self.known_ids)) self.tracking_data["metadata"]["id_mapping"] = {str(k): v for k, v in self.track_to_global.items()} # Save JSON file with open(output_json_path, 'w') as f: json.dump(self.tracking_data, f, indent=2) logger.info(f"Video processing completed. Saved to {output_video_path}") logger.info(f"Video file size: {os.path.getsize(output_video_path)} bytes") return output_video_path, output_json_path # ========== Processor ========== processor = VideoProcessor() # ========== Gradio Handler ========== def process_video_gradio(input_video, progress=gr.Progress()): if input_video is None: return None, None, "Please upload a video file." try: def progress_callback(prog, message): progress(prog, desc=message) # Process video output_video_path, output_json_path = processor.process_video(input_video, progress_callback) # Verify files exist and are accessible if not os.path.exists(output_video_path): raise Exception(f"Output video not found at {output_video_path}") if not os.path.exists(output_json_path): raise Exception(f"Output JSON not found at {output_json_path}") # Read tracking data for stats with open(output_json_path, 'r') as f: data = json.load(f) stats = f""" **Processing Complete!** ✅ - **Total Frames Processed:** {data['metadata']['total_frames']} - **Total People Detected:** {data['metadata']['total_people']} - **Unique IDs Assigned:** {len(data['metadata']['id_mapping'])} - **Output Video Size:** {os.path.getsize(output_video_path) / (1024*1024):.1f} MB 📹 **Output video** is ready for download 📄 **JSON tracking data** contains frame-by-frame detection results """ logger.info(f"Returning video path: {output_video_path}") logger.info(f"Video exists: {os.path.exists(output_video_path)}") return output_video_path, output_json_path, stats except Exception as e: logger.exception("Video processing failed.") return None, None, f"❌ **Error processing video:** {str(e)}" # ========== Gradio Interface ========== def create_interface(): with gr.Blocks(title="Video Person Detection & Tracking", theme=gr.themes.Soft()) as demo: gr.Markdown("# 🎥 Video Person Detection & Tracking with ReID") gr.Markdown("Upload a video to detect and track people using YOLOv8, InsightFace, and ReID models for consistent person identification across frames.") with gr.Row(): with gr.Column(scale=1): input_video = gr.Video( label="📂 Upload Input Video", height=400, interactive=True ) process_btn = gr.Button( "🚀 Process Video", variant="primary", size="lg" ) with gr.Column(scale=1): output_video = gr.Video( label="🎬 Processed Video (with tracking)", height=400, interactive=False, show_download_button=True # Enable download button ) download_json = gr.File( label="📊 Download Tracking Data (JSON)", interactive=False ) with gr.Row(): status_text = gr.Markdown("📤 Upload a video and click **'Process Video'** to start tracking people.") # Event handler process_btn.click( fn=process_video_gradio, inputs=[input_video], outputs=[output_video, download_json, status_text], show_progress=True ) # Additional information with gr.Accordion("📖 How it works", open=False): gr.Markdown(""" ### 🔧 **Technology Stack:** - **YOLOv8:** Real-time person detection - **ByteTrack:** Multi-object tracking algorithm - **InsightFace:** Facial feature extraction for person identification - **OSNet:** Full-body re-identification features ### 📋 **Process:** 1. **Detection:** YOLOv8 detects people in each frame 2. **Tracking:** ByteTrack assigns temporary tracking IDs 3. **Feature Extraction:** InsightFace + OSNet extract identifying features 4. **Re-identification:** Combines face and body features for consistent global IDs 5. **Output:** Generates annotated video + detailed JSON tracking data ### 📁 **Supported Formats:** - **Input:** MP4, AVI, MOV, WEBM - **Output:** MP4 video + JSON metadata """) with gr.Accordion("⚙️ Model Configuration", open=False): gr.Markdown(f""" - **Detection Threshold:** {DETECTION_THRESHOLD} - **Similarity Threshold:** 0.6 (for person re-identification) - **Device:** {"CUDA" if torch.cuda.is_available() else "CPU"} - **Output Directory:** {OUTPUT_DIR} """) with gr.Accordion("🔧 Troubleshooting", open=False): gr.Markdown(""" **If video doesn't display:** 1. Check if the output file exists in the outputs directory 2. Try downloading the video manually 3. Ensure proper video codec support **Common issues:** - Large video files may take time to load - Some browsers may not support certain video formats - Network issues can affect video streaming """) return demo # ========== Launch ========== if __name__ == "__main__": demo = create_interface() # Add file serving for outputs directory demo.launch( server_name="0.0.0.0", server_port=7860, show_error=True, share=True )