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.
    """

    model_path = "racistv4.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(['{:<16}:{:>8}'.format(row[0], row[1]) for row in top5])


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


outputs = gr.Textbox()

title = "Racist model v4"

SAMPLE_DIR = 'samples'
examples = [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)