import gradio as gr
import cv2
import torch
from models import create_YOLO_model
from timeit import default_timer as timer
import numpy as np

# Create the model
model = create_YOLO_model()
model.to('cpu')

def predict(video_path):
    start_time = timer()

    results = model.predict(video_path)

    for result in results:  
        path = result.save("output.mp4")

    pred_time = round(timer() - start_time, 5)

    return path, pred_time

# Define the Gradio interface

# video_example_path = "data/input_video.mp4" 

# example_list = [video_example_path]

iface = gr.Interface(
    fn=predict,
    inputs=gr.Video(label="Upload a Video"),
    outputs=[gr.Video(label="Detection Result"),
             gr.Number(label="Prediction time (s)"),],
    title="YOLO Video Object Detection",
    description="Upload a video and get object detection results using YOLO model."
)

if __name__ == "__main__":
    iface.launch()