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

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


def predict_image(image, conf_threshold, iou_threshold):
    results = pose.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.25,
                  label="Confidence threshold"),
        gr.Slider(minimum=0, maximum=1, value=0.7, label="IoU threshold"),
    ],
    outputs=gr.Image(type="numpy", label="Result"),
    title="Human Pose",
    description="Limbs in all of the right places. 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()