Spaces:
Sleeping
Sleeping
File size: 3,103 Bytes
289a4db f5710ab 92ce950 b763b36 92ce950 b763b36 289a4db b763b36 289a4db b763b36 289a4db b763b36 92ce950 b763b36 92ce950 b763b36 92ce950 289a4db 92ce950 b763b36 92ce950 f5710ab b763b36 92ce950 f5710ab b763b36 289a4db 92ce950 289a4db b763b36 f5710ab b763b36 92ce950 b763b36 92ce950 b763b36 92ce950 b763b36 92ce950 f5710ab 289a4db |
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 |
```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()
``` |