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)