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

classify = YOLO("models/classify.pt")


def predict_image(image, conf_threshold, iou_threshold):
    results = classify.predict(
        image, conf=conf_threshold, iou=iou_threshold, stream=True)

    for r in results:
        im_array = r.plot(labels=True, boxes=True)
        yield Image.fromarray(im_array[..., ::-1])


iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Video(label="Upload Video"),
        gr.Slider(minimum=0, maximum=1, value=0.85,
                  label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.7, label="IoU threshold"),
    ],
    outputs=gr.Image(type="numpy", label="Result"),
    title="Basketball Classifier",
    description="Have you ever wondered where the ball was when you were playing basketball? Where the rim was? Where you were? Videos may take a LOT of time since this is running on the basic CPU tier of HuggingFace. Feel free to check out the image space for a much faster demo!",
)

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