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

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


def predict(model, image, conf_threshold, iou_threshold, show=True):
    results = model.predict(image, conf=conf_threshold, iou=iou_threshold)

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

    return image


def predict_image(image, conf_threshold, iou_threshold):
    return predict(pose, image, conf_threshold, iou_threshold, False)


iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        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="pil", label="Result"),
    title="Human Pose",
    description="Limbs in all of the right places.",
    examples=[
        ["assets/klay.jpeg", 0.85, 0.7],
        ["assets/pierre.png", 0.85, 0.7],
    ]
)

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