import os

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

MODEL_NAME = os.getenv("MODEL_NAME")
HF_TOKEN = os.getenv("HF_TOKEN")

login(token=HF_TOKEN)

model = from_pretrained_keras(MODEL_NAME)


def classify_image(inp):
    inp = cv2.resize(inp, (299, 299))
    inp = np.expand_dims(inp, axis=0)
    prediction = model.predict(inp)
    pred = tf.nn.sigmoid(prediction).numpy().squeeze()
    confidences = {"Application": 1 - pred, "Product": pred.item()}
    return confidences


gr.Interface(
    fn=classify_image,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=2),
    allow_flagging="never",
).launch(debug=True, enable_queue=True)