import os
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO

image_directory = '/home/user/app/example_images'

os.makedirs(image_directory, exist_ok=True)

img_files = [file for file in os.listdir(
    image_directory) if file.lower().endswith('.jpg') or file.lower().endswith('.png')]

path = [os.path.join(image_directory, filename) for filename in img_files]

model = YOLO('/home/user/app/train_cls_best_small.pt')

inputs_image = [
    gr.components.Image(type="filepath", label="Input Image"),
]
outputs_text = [
    gr.components.Textbox(type="text", label="Model predict"),
]


def show_preds_image(image_path):

    results = model(image_path)

    names_dict = results[0].names

    probs = results[0].probs.data.tolist()

    return names_dict[np.argmax(probs)], np.max(probs)


interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_text,
    title="Cat vs Dog",
    examples=path,
    cache_examples=False,
)

gr.TabbedInterface(
    [interface_image],
    tab_names=['Image Inference'],
).queue().launch(debug=True)