import gradio as gr
from ultralytics import YOLO

# catgories
categories = ['Elefante', 'Hipopótamos', 'Rinocerontes']

# returning classifiers output

def image_classifier(inp):
    model = YOLO("best.pt")

    result = model.predict(source=inp)
    probs = result[0].probs.data

    # Combine the two lists and sort based on values in descending order
    sorted_pairs = sorted(zip(categories, probs), key=lambda x: x[1], reverse=True)

    resultado = []
    for name, value in sorted_pairs:
        resultado.append(f'{name}: {value:.2f}%')

    return ', '.join(resultado)

# gradio code block for input and output
with gr.Blocks() as app:
    gr.Markdown("## Classificação de animais de grande porte (Elefantes, Hipopótamos e Rinocerontes) com Yolo-v8")
    with gr.Row():
        inp_img = gr.Image()
        out_txt = gr.Textbox()
    btn = gr.Button(value="Submeter")
    btn.click(image_classifier, inputs=inp_img, outputs=out_txt)

    gr.Markdown("## Exemplos")
    gr.Examples(
        examples=['image_0.jpg', 'image_1.jpg', 'image_2.jpg'],
        inputs=inp_img,
        outputs=out_txt,
        fn=image_classifier,
        cache_examples=True,
    )

app.launch(share=True)