|
import gradio as gr |
|
from ultralytics import YOLO |
|
import tempfile |
|
import cv2 |
|
|
|
|
|
model = "best.pt" |
|
|
|
|
|
def predict_image(image): |
|
results = model.predict(image) |
|
return results[0].plot() |
|
|
|
|
|
def predict_video(video_path): |
|
|
|
cap = cv2.VideoCapture(video_path) |
|
fps = cap.get(cv2.CAP_PROP_FPS) |
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2 |
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2 |
|
|
|
|
|
temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) |
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v') |
|
out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height)) |
|
|
|
while True: |
|
ret, frame = cap.read() |
|
if not ret: |
|
break |
|
|
|
frame = cv2.resize(frame, (width, height)) |
|
results = model.predict(frame, imgsz=480, conf=0.5, verbose=False) |
|
annotated = results[0].plot() |
|
out.write(annotated) |
|
|
|
cap.release() |
|
out.release() |
|
return temp_output.name |
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# π Optimized YOLOv8 Detection\nFast & Accurate on Images and Videos") |
|
|
|
with gr.Tab("Image"): |
|
img_input = gr.Image(type="pil") |
|
img_output = gr.Image(label="Detected") |
|
img_btn = gr.Button("Run Detection") |
|
img_btn.click(predict_image, inputs=img_input, outputs=img_output) |
|
|
|
with gr.Tab("Video"): |
|
vid_input = gr.Video() |
|
vid_output = gr.Video() |
|
vid_btn = gr.Button("Run Detection on Video") |
|
vid_btn.click(predict_video, inputs=vid_input, outputs=vid_output) |
|
|
|
demo.launch() |
|
|