Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import cv2 | |
| import tempfile | |
| from ultralytics import YOLO | |
| from huggingface_hub import hf_hub_download | |
| import spaces | |
| # CONFIGURATION | |
| REPO_ID = "RohanSardar/yolo11s_drone_tracker" | |
| FILENAME = "best.pt" | |
| # Download model from HuggingFace Hub | |
| print(f"Downloading model from {REPO_ID}...") | |
| try: | |
| model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME) | |
| model = YOLO(model_path) | |
| print("Model loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| print("Ensure REPO_ID is set correctly above.") | |
| model = None | |
| try: | |
| bytetrack_path = hf_hub_download(repo_id=REPO_ID, filename="bytetrack.yaml") | |
| print("Custom bytetrack.yaml loaded from HF!") | |
| except Exception: | |
| bytetrack_path = None | |
| print("No custom bytetrack.yaml in HF repo, using default.") | |
| # INFERENCE LOGIC | |
| def process_video(input_video, conf_threshold): | |
| if not model: | |
| return None, "Model not loaded. Please configure REPO_ID in app.py" | |
| if not input_video: | |
| return None, "Please upload a video." | |
| # Setup output path | |
| output_dir = tempfile.mkdtemp() | |
| out_path = os.path.join(output_dir, "tracked_output.mp4") | |
| # Read input video properties | |
| cap = cv2.VideoCapture(input_video) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # We use mp4v; Gradio automatically handles standard mp4 browser playback conversions on the backend | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h)) | |
| # Use custom tracker if available, else built-in | |
| tracker_config = "bytetrack.yaml" | |
| if bytetrack_path and os.path.exists(bytetrack_path): | |
| tracker_config = bytetrack_path | |
| elif os.path.exists("configs/bytetrack.yaml"): | |
| tracker_config = "configs/bytetrack.yaml" | |
| elif os.path.exists("bytetrack.yaml"): | |
| tracker_config = "bytetrack.yaml" | |
| try: | |
| # Run YOLO ByteTrack | |
| results = model.track( | |
| source=input_video, | |
| tracker=tracker_config, | |
| conf=conf_threshold, | |
| stream=True, | |
| verbose=False | |
| ) | |
| for r in results: | |
| # r.plot() returns the BGR image array with boxes drawn | |
| annotated_frame = r.plot() | |
| writer.write(annotated_frame) | |
| except Exception as e: | |
| cap.release() | |
| writer.release() | |
| return None, f"Error processing video: {str(e)}" | |
| cap.release() | |
| writer.release() | |
| # Convert to standard H264 for guaranteed browser compatibility on HF Spaces | |
| final_output = os.path.join(output_dir, "final_h264.mp4") | |
| exit_code = os.system(f"ffmpeg -y -i {out_path} -vcodec libx264 -acodec aac {final_output} 2> NUL") | |
| if exit_code == 0 and os.path.exists(final_output): | |
| return final_output, final_output | |
| print("Note: FFmpeg not found or failed. Returning uncompressed mp4v file.") | |
| return out_path, out_path | |
| # GRADIO INTERFACE | |
| with gr.Blocks() as app: | |
| gr.Markdown( | |
| """ | |
| # Drone Tracker | |
| Upload a video to detect and track small drones. | |
| Powered by fine-tuned YOLOv11 and ByteTrack. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_input = gr.Video(label="Upload Video") | |
| conf_slider = gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="Confidence Threshold") | |
| submit_btn = gr.Button("Start Tracking", variant="primary") | |
| with gr.Column(): | |
| video_output = gr.Video(label="Tracked Output", interactive=False) | |
| download_file = gr.File(label="Download File") | |
| error_text = gr.Markdown() | |
| submit_btn.click( | |
| fn=process_video, | |
| inputs=[video_input, conf_slider], | |
| outputs=[video_output, download_file] | |
| ) | |
| if __name__ == "__main__": | |
| app.launch(theme=gr.themes.Soft()) | |