import gradio as gr
from tensorflow.keras.models import load_model
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image

# --- 1. Load the Model from your other Hugging Face Repo ---
try:
    model_path = hf_hub_download(
        repo_id="skibi11/leukolook-eye-detector", 
        filename="MobileNetV1_best.keras"
    )
    model = load_model(model_path)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

# --- 2. Define the Pre-processing Logic ---
def preprocess_image(img_pil):
    # This MUST match your training pre-processing
    img = img_pil.resize((224, 224))
    img_array = np.array(img)
    if img_array.ndim == 2:
        img_array = np.stack((img_array,)*3, axis=-1)
    img_array = img_array / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# --- 3. Define the Prediction Function ---
def predict(image_array):
    if model is None:
        raise gr.Error("Model is not loaded. Please check the Space logs.")

    pil_image = Image.fromarray(image_array.astype('uint8'), 'RGB')
    processed_image = preprocess_image(pil_image)
    prediction = model.predict(processed_image)

    # Convert prediction to a JSON-friendly format
    labels = [f"Class_{i}" for i in range(prediction.shape[1])]
    confidences = {label: float(score) for label, score in zip(labels, prediction[0])}
    return confidences

# --- 4. Create and Launch the Gradio API ---
gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs="json",
    title="LeukoLook Eye Detector API"
).launch()