import gradio as gr
import torch
from ultralyticsplus import YOLO, render_result


# Images
torch.hub.download_url_to_file('https://external-content.duckduckgo.com/iu/?u=https%3A%2F%2Ftexashafts.com%2Fwp-content%2Fuploads%2F2016%2F04%2Fconstruction-worker.jpg', 'one.jpg')
torch.hub.download_url_to_file(
    'https://www.pearsonkoutcherlaw.com/wp-content/uploads/2020/06/Construction-Workers.jpg', 'two.jpg')
torch.hub.download_url_to_file(
    'https://nssgroup.com/wp-content/uploads/2019/02/Building-maintenance-blog.jpg', 'three.jpg')


def yoloV8_func(image: gr.inputs.Image = None,
                image_size: gr.inputs.Slider = 640,
                conf_threshold: gr.inputs.Slider = 0.4,
                iou_threshold: gr.inputs.Slider = 0.50):
    """_summary_
    Args:
        image (gr.inputs.Image, optional): _description_. Defaults to None.
        image_size (gr.inputs.Slider, optional): _description_. Defaults to 640.
        conf_threshold (gr.inputs.Slider, optional): _description_. Defaults to 0.4.
        iou_threshold (gr.inputs.Slider, optional): _description_. Defaults to 0.50.
    """
    model_path = "best.pt"
    model = YOLO("foduucom/table-detection-and-extraction")

    results = model.predict(image,
                            conf=conf_threshold,
                            iou=iou_threshold,
                            imgsz=image_size)

    # observe results
    box = results[0].boxes
    print("Object type:", box.cls)
    print("Coordinates:", box.xyxy)
    print("Probability:", box.conf)
    render = render_result(model=model, image=image, result=results[0])
    return render


inputs = [
    gr.inputs.Image(type="filepath", label="Input Image"),
    gr.inputs.Slider(minimum=320, maximum=1280, default=640,
                     step=32, label="Image Size"),
    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25,
                     step=0.05, label="Confidence Threshold"),
    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45,
                     step=0.05, label="IOU Threshold"),
]


outputs = gr.outputs.Image(type="filepath", label="Output Image")
title = "YOLOv8 101: Custome Object Detection on Construction Workers "


examples = [['one.jpg', 640, 0.5, 0.7],
            ['two.jpg', 800, 0.5, 0.6],
            ['three.jpg',  900, 0.5, 0.8]]

yolo_app = gr.Interface(
    fn=yoloV8_func,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    cache_examples=True,
    #theme='huggingface',
)
yolo_app.launch(debug=True, enable_queue=True)