import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras import backend as K

# ---- Custom metrics/loss required to load your model ---- #
def dice_coef(y_true, y_pred):
    y_true = K.cast(y_true, 'float32')
    y_pred = K.cast(y_pred, 'float32')
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersect = K.sum(y_true_f * y_pred_f)
    return (2. * intersect + K.epsilon()) / (K.sum(y_true_f) + K.sum(y_pred_f) + K.epsilon())

# ---- Load model from Hugging Face model repo ---- #
model_path = "yaraa11/brain-tumor-segmentation"   # <-- Replace with your HF model ID
model = load_model(model_path, custom_objects={'dice_coef': dice_coef}, compile=False)

IMG_SIZE = (224, 224)

# ---- Prediction function for Gradio ---- #
def predict(img):
    # Convert to grayscale if not already
    img = img.convert("L")

    # Resize
    img_resized = img.resize(IMG_SIZE)
    x = np.array(img_resized) / 255.0
    x = np.expand_dims(x, axis=(0, -1))  # shape: (1, 224, 224, 1)

    # Predict mask
    pred = model.predict(x)[0]
    mask = (pred > 0.5).astype(np.uint8) * 255  # Convert to binary mask

    # Convert numpy mask to image
    mask_img = Image.fromarray(mask.squeeze(), mode="L")
    return mask_img

# ---- Gradio Interface ---- #
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Brain MRI Image"),
    outputs=gr.Image(type="pil", label="Predicted Tumor Mask"),
    title="Brain Tumor Segmentation",
    description="Upload a brain MRI image to generate a tumor segmentation mask."
)

interface.launch()