import os
import time

import cv2
import gradio as gr
from PIL import Image
from ultralytics import YOLO

model = YOLO("yolo26_75ep_640_drone_detector_openvino_model")
example_list = [["examples/" + example] for example in os.listdir("examples")]


def predict_video_stream(video_path, conf_threshold, iou_threshold):
    results = model.track(
        source=video_path,
        conf=conf_threshold,
        iou=iou_threshold,
        persist=True,
        stream=True,
        save=False,
        vid_stride=2,
    )

    for frame_results in results:
        annotated_frame = frame_results.plot()

        rgb_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)

        small_frame = cv2.resize(rgb_frame, (640, 480))

        pil_img = Image.fromarray(small_frame)

        yield pil_img

        time.sleep(0.02)


with gr.Blocks() as demo:
    with gr.Row():
        input_video = gr.Video(label="Upload video")
        output_image = gr.Image(label="Tracking in real-time", type="numpy")

    btn = gr.Button("RUN TRACKING")

    btn.click(
        fn=predict_video_stream,
        inputs=[input_video, gr.Slider(0, 1, value=0.25), gr.Slider(0, 1, value=0.45)],
        outputs=output_image,
    )

    gr.Examples(
        examples=example_list,
        inputs=input_video
    )

demo.queue().launch()