Spaces:
Sleeping
Sleeping
File size: 4,016 Bytes
b9703d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | 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
@spaces.GPU
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())
|