Spaces:
Sleeping
Sleeping
```python | |
import os | |
import gradio as gr | |
import cv2 | |
import numpy as np | |
from ultralytics import YOLO | |
import tempfile | |
# Set Ultralytics config path | |
os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics' | |
# Custom function to load model with weights_only=False | |
def load_model_with_trusted_weights(model_path): | |
import torch | |
with torch.serialization.safe_globals(): | |
return YOLO(model_path) | |
# Load both YOLO models | |
model_yolo11 = load_model_with_trusted_weights('./data/yolo11n.pt') | |
model_best = load_model_with_trusted_weights('./data/best.pt') | |
def process_video(video_path, model_name, conf_threshold=0.4): | |
""" | |
Process the input video frame by frame using the selected YOLO model, | |
draw bounding boxes, and return the processed video path. | |
""" | |
# Select model | |
model = model_yolo11 if model_name == "YOLO11n" else model_best | |
# Open video capture | |
cap = cv2.VideoCapture(video_path) | |
if not cap.isOpened(): | |
raise ValueError("Could not open video file") | |
# Get video properties | |
fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
# Define output video path | |
temp_video_path = os.path.join(tempfile.gettempdir(), "output.mp4") | |
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for MP4 | |
out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height)) | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
# Perform detection | |
results = model.predict( | |
source=frame, | |
conf=conf_threshold, | |
imgsz=640, | |
show_labels=True, | |
show_conf=True | |
) | |
# Draw bounding boxes | |
for result in results: | |
im_array = result.plot() # Plot boxes | |
out.write(im_array) # Write frame to output video | |
cap.release() | |
out.release() | |
cv2.destroyAllWindows() | |
return temp_video_path | |
# Gradio interface | |
with gr.Blocks() as app: | |
gr.HTML(""" | |
<h1 style='text-align: center'> | |
Video Object Detection with YOLO Models | |
</h1> | |
""") | |
with gr.Row(): | |
with gr.Column(): | |
video_input = gr.Video(label="Upload Video") | |
model_choice = gr.Dropdown( | |
choices=["YOLO11n", "Best Model"], | |
label="Select Model", | |
value="YOLO11n" | |
) | |
conf_threshold = gr.Slider( | |
label="Confidence Threshold", | |
minimum=0.0, | |
maximum=1.0, | |
step=0.05, | |
value=0.4 | |
) | |
process_button = gr.Button("Process Video") | |
with gr.Column(): | |
video_output = gr.Video( | |
label="Processed Video", | |
streaming=True, | |
autoplay=True | |
) | |
process_button.click( | |
fn=process_video, | |
inputs=[video_input, model_choice, conf_threshold], | |
outputs=[video_output] | |
) | |
if __name__ == "__main__": | |
app.launch() | |
``` |