Yolo-IFMT / app.py
LimaLimao's picture
Upload 4 files
0f65929 verified
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,
# conf=conf_threshold,
# iou=iou_threshold,
# show_labels=True,
# show_conf=True,
)
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='./')
# print(results)
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)