import tensorflow as tf
import gradio as gr
import numpy as np
import json
from PIL import Image
from tensorflow.keras.applications.efficientnet import preprocess_input

# Load model
model = tf.keras.models.load_model("polyp_efficientnet_model.h5")

# Load class names
with open("class_names.json") as f:
    class_names = json.load(f)

IMG_SIZE = (224, 224)

def predict(image):
    if image is None:
        return None

    # Resize
    image = image.resize(IMG_SIZE)

    # Convert to numpy
    image = np.array(image)

    # Ensure RGB
    if image.shape[-1] == 4:
        image = image[..., :3]

    # Batch dimension
    image = np.expand_dims(image, axis=0)

    # ✅ CORRECT preprocessing
    image = preprocess_input(image)

    preds = model.predict(image)[0]

    return {
        class_names[i]: float(preds[i])
        for i in range(len(class_names))
    }

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=4),
    title="Polyp Disease Classification",
    description="EfficientNet-based medical image classifier"
)

interface.launch()