# Final robust app.py for your Hugging Face Space

import gradio as gr
import tensorflow as tf
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import os

# --- 1. Load the Model ---
model = None
try:
    model_path = hf_hub_download(
        repo_id="skibi11/leukolook-eye-detector", 
        filename="MobileNetV1_best.keras"
    )
    model = tf.keras.models.load_model(model_path)
    print("--- MODEL LOADED SUCCESSFULLY! ---")
except Exception as e:
    print(f"--- ERROR LOADING MODEL: {e} ---")
    raise gr.Error(f"Failed to load model: {e}")

# --- 2. Pre-processing Logic ---
def preprocess_image(img_pil):
    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)
    if img_array.shape[-1] == 4:
        img_array = img_array[..., :3]
    img_array = img_array / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# --- 3. Prediction Logic ---
def predict(image_from_gradio):
    if not isinstance(image_from_gradio, np.ndarray):
        return {"error": "Invalid input type. Expected an image."}

    try:
        pil_image = Image.fromarray(image_from_gradio)
        processed_image = preprocess_image(pil_image)
        prediction = model.predict(processed_image)

        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
    except Exception as e:
        raise gr.Error(f"Error during prediction: {e}")

# --- 4. Create Gradio Interface using gr.Blocks for stability ---
with gr.Blocks() as demo:
    gr.Interface(
        fn=predict,
        inputs=gr.Image(type="numpy"),
        outputs=gr.JSON(),
        api_name="predict"
    )

demo.launch()