import gradio as gr
import torch
from ultralytics import YOLO

torch.hub.download_url_to_file(
    'http://www.imesclub.org/images/stories/arabidentity.jpeg', 'one.jpg')
torch.hub.download_url_to_file(
    'https://lp-cms-production.imgix.net/2023-08/GettyImages-1224548888-16.9.jpg', 'two.jpg')
torch.hub.download_url_to_file(
    'https://s.wsj.net/public/resources/images/OB-EQ974_diwali_G_20091014112400.jpg', 'three.jpg')


def handle_classify(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.
    """
    
    if not image:
        return "No image found"

    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 "\n".join(["\t".join(row) for row in top5])


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


outputs = gr.Textbox()

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)