import os
import gradio as gr
from huggingface_hub import login
from huggingface_hub import snapshot_download
import numpy as np
import tensorflow as tf
import cv2

def image_mod(image):
    # Convert image from PIL format (RGBA) to the cv2 format
    # (BGRA) that was used when training the model
    res = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
    res = image.astype("float32") / 255.0

    # Convert single image to a batch for prediction
    res = np.array([res])

    # Perform prediction
    result = 1.0 - model.predict(res)[0][0]

    # Return the result
    return result

# Download the model from Hugging Face Hub
login(token=os.environ['TOKEN_KESTREL_READ'])
model_path = snapshot_download(repo_id=os.environ['REPO_KESTREL_MODEL'])
model = tf.keras.models.load_model(model_path)

# Configure Gradio components
input_image_component = gr.Image(image_mode="RGBA", type="numpy")
output_number_component = gr.Number()

# Configure user interface
iface = gr.Interface(
    fn=image_mod,
    inputs=input_image_component,
    outputs=output_number_component,
    live=True,
    title="",
    description="",
    article=""
)

# Launch the frontend server
iface.launch(
    share=False,
    auth=(os.environ['USER_KESTREL'],
    os.environ['PASSWORD_KESTREL'])
)