# Final app.py for your Hugging Face Space

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

# --- 1. Load the Model from your other Hugging Face Repo ---
model = None
print("--- SCRIPT START ---")
try:
    print("Downloading Keras model from the Hub...")
    model_path = hf_hub_download(
        repo_id="skibi11/leukolook-eye-detector", 
        filename="MobileNetV1_best.keras"
    )
    print(f"Model downloaded to: {model_path}")
    print("Loading model with tf.keras.models.load_model...")

    # This is a more robust way to load the model
    model = tf.keras.models.load_model(model_path)

    print("--- MODEL LOADED SUCCESSFULLY! ---")
    model.summary() # Print a summary of the model to confirm it's loaded

except Exception as e:
    print("--- AN ERROR OCCURRED DURING MODEL LOADING ---")
    print(f"Error Type: {type(e)}")
    print(f"Error Message: {e}")
    # Also print the traceback for more details
    import traceback
    traceback.print_exc()
    print("--- END OF ERROR ---")


# --- 2. Define the 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)
    # Ensure image has 3 channels if it's not
    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. Define the Prediction Function ---
def predict(image_array):
    if model is None:
        raise gr.Error("Model is not loaded. Please check the Space logs for errors.")

    try:
        pil_image = Image.fromarray(image_array.astype('uint8'), 'RGB')
        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 and Launch the Gradio API ---
gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs="json",
    title="LeukoLook Eye Detector API",
    description="API for the LeukoLook project."
).launch()