import gradio as gr
from huggingface_hub import from_pretrained_keras
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.preprocessing import image
import numpy as np

model = from_pretrained_keras("yusyel/fishv2")


class_names = [
    "Black Sea Sprat",
    "Gilt-Head Bream",
    "Hourse Mackerel",
    "Red Sea Bream",
    "Red Mullet",
    "Sea Bass",
    "Shrimp",
    "Striped Red Mullet",
    "Trout",
]


def preprocess_image(img, label):
    img = load_img(img, target_size=(199, 199))
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img /= 255.0
    print(img.shape)
    return img, label



def predict(img):
    img, _ = preprocess_image(img, 1)
    pred = model.predict(img)
    pred = np.squeeze(pred).astype(float)
    print(pred)
    return dict(zip(class_names, pred))


demo = gr.Interface(
    fn=predict,
    inputs=[gr.inputs.Image(type="filepath")],
    outputs=gr.outputs.Label(),
    examples=[
        ["./img/Black_Sea_Sprat.png"],
        ["./img/Gilt_Head_Bream.JPG"],
        ["./img/Horse_Mackerel.png"],
        ["./img/Red_mullet.png"],
        ["./img/Red_Sea_Bream.JPG"],
        ["./img/Sea_Bass.JPG"],
        ["./img/Shrimp.png"],
        ["./img/Striped_Red_Mullet.png"],
        ["./img/Trout.png"],
    ],
    title="fish classification",
)

demo.launch()