hf_spaces_app.py

"""
HuggingFace Spaces entry point — Gradio demo for SwinUNETR.
Accepts NIfTI uploads, returns axial/coronal/sagittal slice PNGs.
"""
import os
import sys
import tempfile
import uuid
from pathlib import Path

import numpy as np
import torch
import yaml
import nibabel as nib
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

from src.model import create_model
from src.utils import get_inferer, load_model
from monai.transforms import Activations, AsDiscrete

# ── Config ────────────────────────────────────────────────────────────────────
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CONFIG_PATH     = "config/config.yaml"
CHECKPOINT_PATH = Path("runs/best_fold0.pt")

cfg = yaml.safe_load(open(CONFIG_PATH)) if os.path.exists(CONFIG_PATH) else {
    "in_channels": 4, "out_channels": 4, "feature_size": 48,
    "roi": [128, 128, 128], "sw_batch_size": 1, "infer_overlap": 0.25,
}

# ── Load Model ────────────────────────────────────────────────────────────────
model = None
if CHECKPOINT_PATH.exists():
    model = load_model(str(CHECKPOINT_PATH), cfg, DEVICE)
    print(f"[HF Space] Model loaded on {DEVICE}")
else:
    print("[HF Space] No checkpoint — demo will show error")

# ── Colors ────────────────────────────────────────────────────────────────────
COLORS = {1: [1,0,0], 2: [0,1,0], 3: [0,0,1]}  # TC=red, WT=green, ET=blue


def nifti_to_tensor(path: str) -> torch.Tensor:
    img = nib.load(path).get_fdata(dtype=np.float32)
    if img.ndim == 3:
        img = np.expand_dims(img, 0)
    elif img.ndim == 4:
        img = np.transpose(img, (3, 0, 1, 2))
    for c in range(img.shape[0]):
        ch = img[c]
        nz = ch[ch != 0]
        if len(nz) > 0:
            ch = (ch - nz.mean()) / (nz.std() + 1e-8)
            img[c] = ch
    return torch.from_numpy(img).float().unsqueeze(0)


def overlay_slice(img_sl, seg_sl, alpha=0.4):
    img_norm = (img_sl - img_sl.min()) / (img_sl.max() - img_sl.min() + 1e-8)
    rgb = np.stack([img_norm, img_norm, img_norm], axis=-1)
    for cid, color in COLORS.items():
        mask = (seg_sl == cid)
        for ch, val in enumerate(color):
            rgb[..., ch] = np.where(mask, rgb[..., ch] * (1-alpha) + val*alpha, rgb[..., ch])
    return np.clip(rgb, 0, 1)


@torch.no_grad()
def predict(nifti_file):
    if nifti_file is None:
        return None, None, None, "Please upload a NIfTI file."

    if model is None:
        return None, None, None, "❌ No model checkpoint found. Train the model first."

    try:
        tensor  = nifti_to_tensor(nifti_file.name).to(DEVICE)
        inferer = get_inferer(model, tuple(cfg["roi"]), cfg["sw_batch_size"], cfg["infer_overlap"])

        post_softmax = Activations(softmax=True)
        post_pred    = AsDiscrete(argmax=True)

        logits  = inferer(tensor)
        pred    = post_pred(post_softmax(logits)).squeeze().cpu().numpy()
        img_np  = tensor.squeeze().cpu().numpy()  # (C, H, W, D)
        flair   = img_np[0]
        H, W, D = flair.shape

        def make_fig(img_sl, seg_sl, title):
            fig, axes = plt.subplots(1, 2, figsize=(10, 4))
            axes[0].imshow(img_sl.T, cmap="gray", origin="lower")
            axes[0].set_title("FLAIR Input")
            axes[0].axis("off")
            axes[1].imshow(overlay_slice(img_sl, seg_sl).transpose(1,0,2), origin="lower")
            axes[1].set_title("Segmentation")
            axes[1].axis("off")
            fig.suptitle(title, fontsize=11)
            fig.tight_layout()
            return fig

        fig_ax  = make_fig(flair[:,:,D//2],  pred[:,:,D//2],  "Axial View")
        fig_cor = make_fig(flair[:,W//2,:],   pred[:,W//2,:],  "Coronal View")
        fig_sag = make_fig(flair[H//2,:,:],   pred[H//2,:,:],  "Sagittal View")

        labels  = np.unique(pred).tolist()
        info    = f"✅ Done | Shape: {pred.shape} | Labels: {[int(l) for l in labels]}"
        return fig_ax, fig_cor, fig_sag, info

    except Exception as e:
        return None, None, None, f"❌ Error: {str(e)}"


# ── Gradio Interface ──────────────────────────────────────────────────────────
with gr.Blocks(title="🧠 Brain Tumor Segmentation") as demo:
    gr.Markdown("""
    # 🧠 3D Brain Tumor Segmentation
    **SwinUNETR** trained on BraTS 2021 — upload a NIfTI MRI scan to get a 3D segmentation.

    Tumor regions:
    🔴 **Tumor Core (TC)** · 🟢 **Whole Tumor (WT)** · 🔵 **Enhancing Tumor (ET)**
    """)

    with gr.Row():
        with gr.Column(scale=1):
            nifti_input = gr.File(
                label="Upload NIfTI file (.nii or .nii.gz)",
                file_types=[".nii", ".gz"],
            )
            run_btn = gr.Button("🧠 Run Segmentation", variant="primary")
            info_box = gr.Textbox(label="Status", interactive=False)

        with gr.Column(scale=3):
            axial_out   = gr.Plot(label="Axial View")
            coronal_out = gr.Plot(label="Coronal View")
            sagittal_out = gr.Plot(label="Sagittal View")

    run_btn.click(
        fn=predict,
        inputs=[nifti_input],
        outputs=[axial_out, coronal_out, sagittal_out, info_box],
    )

    gr.Markdown("""
    ---
    **Model:** SwinUNETR (MONAI) · **Dataset:** BraTS 2021 · **Author:** Moez Bouassida
    """)

if __name__ == "__main__":
    demo.launch()