import gradio as gr
import torch
from ultralytics import YOLO

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 handle_classify(image: gr.inputs.Image = None):
    """This function performs YOLOv8 object detection on the given image.

    Args:
        image (gr.inputs.Image, optional): Input image to detect objects on. Defaults to None.
    """

    model_path = "racist2.0.pt"
    model = YOLO(model_path)

    results = model(image)
    
    result = results[0]
    
    top5 = [[result.names[class_index], result.probs.top5conf.tolist()[rank]]
                    for class_index, rank in zip(result.probs.top5, range(5))]
    
    print(top5)
    
    return top5


inputs = [
    gr.inputs.Image(type="filepath", label="Input Image"),
]


outputs = gr.outputs.Text()

title = "Racist model v2"


examples = [['one.jpg'],
            ['two.jpg'],
            ['three.jpg']]

yolo_app = gr.Interface(
    fn=handle_classify,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    cache_examples=True,
)

# Launch the Gradio interface in debug mode with queue enabled
yolo_app.launch(debug=True, enable_queue=True)