Spaces:
Sleeping
Sleeping
| 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 | |
| ) |