import gradio as gr
from huggingface_hub import login, from_pretrained_keras
import os
import numpy as np
import cv2

login(os.environ["HF_TOKEN"])

modelv1 = from_pretrained_keras("elsamueldev/cats-dogs")
modelv2 = from_pretrained_keras("elsamueldev/cats-dogs-v2")


def preprocess(img: np.array) -> np.array:
    img = cv2.resize(img, (100, 100)) # resize to 100x100
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # convert to grayscale
    img = img.reshape(100, 100, 1) # reshape to 100x100x1
    img = img / 255 # normalize
    img = np.array([img]) # reshape to 1x100x100x1

    return img

def predict(img: np.array, model: str):
    img = preprocess(img)

    if model == "v1":
        dog = modelv1.predict(img)[0][0]
    else:
        dog = modelv2.predict(img)[0][0]
    cat = 1 - dog

    return {"dog": dog, "cat": cat}

gr.Interface(
    fn=predict,
    inputs=["image", gr.Dropdown(choices=["v1", "v2"], value="v2")],
    outputs="label"
).launch()