import gradio as gr
from PIL import Image
import numpy as np
from ultralytics import YOLO
import os

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], str(round(result.probs.top5conf.tolist()[rank], 4)*100)+'%']
                    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='numpy', label="Input Image"),
]


outputs = gr.Textbox()

title = "Racist model v2"

SAMPLE_DIR = 'samples'
examples = [np.array(Image.open(os.path.join(SAMPLE_DIR, path))) for path in os.listdir(SAMPLE_DIR)]

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)