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

# Load models from your repos
classification_model_path = hf_hub_download(
    repo_id="MohammedAH/Brrain-MRI-Classification",
    filename="brain_mri.h5"
)

segmentation_model_path = hf_hub_download(
    repo_id="MohammedAH/Unet-Brain-Segmentation",
    filename="Unet_model.h5"
)

classification_model = load_model(classification_model_path, compile=False)
segmentation_model = load_model(segmentation_model_path, compile=False)

class_names = ['glioma', 'meningioma', 'no_tumor', 'pituitary']


def predict(image):
    # Classification
    img = cv2.resize(image, (224, 224)) / 255.0
    cls_input = np.expand_dims(img, axis=0)

    preds = classification_model.predict(cls_input)
    idx = int(np.argmax(preds[0]))

    # Segmentation
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    seg = cv2.resize(gray, (128, 128)) / 255.0
    seg_input = np.expand_dims(seg, axis=(0, -1))

    mask = segmentation_model.predict(seg_input)
    mask = (mask > 0.5).astype(np.uint8)[0, :, :, 0]

    return {
        "prediction": class_names[idx],
        "confidence": float(preds[0][idx]),
        "mask": mask.tolist()
    }


interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.JSON()   # 🔥 important (API-friendly)
)

interface.launch()