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

# 🔥 Download models from your Hugging Face 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"
)

# Load models once
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 preprocessing
    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 preprocessing
    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 class_names[idx], float(preds[0][idx]), mask


interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Text(label="Prediction"),
        gr.Number(label="Confidence"),
        gr.Image(label="Segmentation Mask")
    ]
)

interface.launch()