import gradio as gr
import requests
from PIL import Image
import io

from Model import FishModel

Classifier =  FishModel()

def inference(image_input, url_input) -> dict[str, float]:
    image = load_image(image_input, url_input)

    result = Classifier.get_fish_species(image)

    data = dict()

    for element in result:
        data[element['label']] = round(element['score'], 2)

    return data # {"dog" : 30, "cat" : 70}

def load_image(image_input, url_input):
    if image_input is not None:
        return image_input
        
    elif url_input:
        try:
            response = requests.get(url_input)
            response.raise_for_status()
            return Image.open(io.BytesIO(response.content))
        
        except Exception as e:
            raise gr.Error(f"No se pudo cargar la imagen desde la URL. Error: {e}")


demo = gr.Interface(
    title="🐳📸 Fish Classification",
    description=open("description.md", "r", encoding="utf8").read(),
    fn=inference,
    inputs=[
        gr.Image(label="Sube una imagen", type="pil"),
        gr.Textbox(label="O pega una URL de imagen aquí")
    ],
    outputs=gr.Label(label="Resultado"),
    # examples=[
    #     [None, "https://gradio-builds.s3.amazonaws.com/demo-files/goldfish.jpg"],
    #     ["images/salmon.jpg", None]
    # ]
)

demo.launch()