|
|
import gradio as gr
|
|
|
import cv2
|
|
|
from ultralytics import YOLOv10
|
|
|
from moviepy.editor import *
|
|
|
|
|
|
|
|
|
model = YOLOv10("best.pt")
|
|
|
|
|
|
def predict_image(img):
|
|
|
results = model.predict(
|
|
|
source=img,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
)
|
|
|
im_rgb = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
|
|
|
return im_rgb if results else None
|
|
|
|
|
|
|
|
|
def predict_video(video):
|
|
|
results = model.predict(source=video, save_dir='./')
|
|
|
|
|
|
images_list = []
|
|
|
|
|
|
for r in results:
|
|
|
im_rgb = cv2.cvtColor(r.plot(), cv2.COLOR_BGR2RGB)
|
|
|
images_list.append(im_rgb)
|
|
|
|
|
|
clip = ImageSequenceClip(images_list, fps=15)
|
|
|
clip.ipython_display(width = 360)
|
|
|
|
|
|
return '__temp__.mp4' if results else None
|
|
|
|
|
|
|
|
|
Image = gr.Interface(fn=predict_image,
|
|
|
inputs=gr.Image(type="pil", label="Upload Image"),
|
|
|
outputs=gr.Image(type="pil", label="Result"),
|
|
|
description="Upload images for YOLOv10 object detection.",
|
|
|
)
|
|
|
|
|
|
Video = gr.Interface(fn=predict_video,
|
|
|
inputs=gr.Video(),
|
|
|
outputs=gr.Video(),
|
|
|
description="Upload Video for YOLOv10 object detection.",
|
|
|
)
|
|
|
|
|
|
demo = gr.TabbedInterface(
|
|
|
[Image, Video],
|
|
|
["Image", "Video"],
|
|
|
title="Ultralytics Gradio YOLOv10",
|
|
|
)
|
|
|
|
|
|
|
|
|
demo.launch(share=True) |