Initial commit with only essential weights

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,20 @@
+FROM python:3.13.5-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt ./
+COPY src/ ./src/
+
+RUN pip3 install -r requirements.txt
+
+EXPOSE 8501
+
+HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
+
+ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]

README.md

@@ -0,0 +1,19 @@
+---
+title: Brain Tumor AI Suite
+emoji: 🚀
+colorFrom: red
+colorTo: red
+sdk: docker
+app_port: 8501
+tags:
+- streamlit
+pinned: false
+short_description: Streamlit template space
+---
+
+# Welcome to Streamlit!
+
+Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
+
+If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
+forums](https://discuss.streamlit.io).

app.py

@@ -0,0 +1,36 @@
+
+import streamlit as st
+
+st.set_page_config(
+    page_title="Brain Tumor AI Suite",
+    page_icon="🧠",
+    layout="wide",
+)
+
+st.markdown('''
+<style>
+    .hero-title {
+        font-size: 3rem; font-weight: 800;
+        background: linear-gradient(135deg, #667eea, #764ba2, #f093fb);
+        -webkit-background-clip: text; -webkit-text-fill-color: transparent;
+        margin-bottom: 0;
+    }
+</style>
+''', unsafe_allow_html=True)
+
+st.markdown('<p class="hero-title">🧠 Brain Tumor AI Suite</p>', unsafe_allow_html=True)
+st.markdown("**Federated Learning Classification & 3D Segmentation**")
+st.markdown("---")
+
+col1, col2 = st.columns(2)
+
+with col1:
+    st.markdown("### 📊 Tumor Classification (Federated Learning)")
+    st.markdown("Predict the class of a brain tumor (Glioma, Meningioma, Pituitary, No Tumor) using a SimpleCNN model trained across simulated hospitals with Layer-by-Layer QPSO aggregation.")
+    st.page_link("pages/1_Classification.py", label="Open Classification App →", icon="📊")
+
+with col2:
+    st.markdown("### 🔬 3D Tumor Segmentation")
+    st.markdown("View MRI slices and 3D volumetric renderings of brain tumors with segmentation overlays (Whole Tumor, Tumor Core, Enhancing Tumor) predicted by a 3D Attention U-Net.")
+    st.page_link("pages/2_Slice_Viewer.py", label="Open Slice Viewer →", icon="🔬")
+    st.page_link("pages/3_3D_Visualization.py", label="Open 3D Viewer →", icon="🌐")

best_metric_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4980fada1cc6fed5b19cd657528845f2a8598b383dcbbba179c1654b6f592c02
+size 23731355

pages/1_Classification.py

@@ -0,0 +1,324 @@
+"""
+🧠 Brain Tumor Classification — Federated Learning Demo
+=========================================================
+Demonstrates the FL-trained SimpleCNN models (FedAvg, FedProx, QPSO).
+Users can upload images or use sample test images from the dataset.
+Shows predicted class with confidence bars and compares all 3 models.
+"""
+
+import streamlit as st
+import os
+import sys
+import glob
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+import plotly.graph_objects as go
+
+
+
+# ─── paths ───────────────────────────────────────────────────────────────
+FL_ROOT = os.path.abspath(os.path.dirname(__file__))
+RESULTS = os.path.join(FL_ROOT, "..", "results")
+# Use Setup 1 models by default (best performing)
+MODELS_DIR = os.path.join(RESULTS, "Setup_1", "models")
+
+NUM_CLASSES = 4
+CLASS_NAMES = ["Glioma", "Meningioma", "No Tumor", "Pituitary"]
+CLASS_COLORS = ["#E74C3C", "#3498DB", "#2ECC71", "#9B59B6"]
+CLASS_ICONS = ["🔴", "🔵", "🟢", "🟣"]
+IMG_SIZE = 112
+
+# ─── model ───────────────────────────────────────────────────────────────
+class SimpleCNN(nn.Module):
+    def __init__(self, num_classes=4):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, padding=1),
+            nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(16, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64), nn.ReLU(),
+            nn.AdaptiveAvgPool2d(4),
+        )
+        self.classifier = nn.Sequential(
+            nn.Dropout(0.3),
+            nn.Linear(64 * 4 * 4, 128), nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(128, num_classes),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        return self.classifier(x)
+
+
+TRANSFORM = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+
+@st.cache_resource
+def load_model(path):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = SimpleCNN(NUM_CLASSES).to(device)
+    state = torch.load(path, map_location=device, weights_only=True)
+    model.load_state_dict(state)
+    model.eval()
+    return model, device
+
+
+def predict(model, device, image):
+    """Run inference on a PIL Image. Returns (class_idx, probabilities)."""
+    tensor = TRANSFORM(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = model(tensor)
+        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
+    return int(np.argmax(probs)), probs
+
+
+def render_prediction_card(title, color_accent, pred_idx, probs, image):
+    """Render a styled prediction result."""
+    confidence = probs[pred_idx] * 100
+    st.markdown(
+        f'<div style="background:rgba(20,20,35,0.9);padding:20px;border-radius:12px;'
+        f'border-top:4px solid {color_accent};margin-bottom:16px;">'
+        f'<h3 style="color:{color_accent};margin:0 0 8px 0;">{title}</h3>'
+        f'<div style="color:white;font-size:28px;font-weight:800;margin:4px 0;">'
+        f'{CLASS_ICONS[pred_idx]} {CLASS_NAMES[pred_idx]}</div>'
+        f'<div style="color:#aaa;font-size:14px;">Confidence: {confidence:.1f}%</div>'
+        f'</div>',
+        unsafe_allow_html=True,
+    )
+
+    # Probability bar chart
+    fig = go.Figure(go.Bar(
+        x=probs * 100,
+        y=CLASS_NAMES,
+        orientation='h',
+        marker_color=CLASS_COLORS,
+        text=[f"{p*100:.1f}%" for p in probs],
+        textposition='auto',
+    ))
+    fig.update_layout(
+        height=200,
+        margin=dict(l=0, r=10, t=10, b=10),
+        paper_bgcolor="rgba(0,0,0,0)",
+        plot_bgcolor="rgba(20,20,35,0.5)",
+        xaxis=dict(range=[0, 100], title="Probability (%)", color="#aaa",
+                   gridcolor="rgba(100,100,140,0.2)"),
+        yaxis=dict(color="white"),
+        font=dict(color="white"),
+    )
+    st.plotly_chart(fig, use_container_width=True)
+
+
+# ─── CSS ─────────────────────────────────────────────────────────────────
+st.markdown("""
+<style>
+    [data-testid="stSidebar"] {
+        background: linear-gradient(180deg, #0f0f1a 0%, #1a1a2e 100%);
+    }
+    .hero-title {
+        font-size: 2.4rem; font-weight: 800;
+        background: linear-gradient(135deg, #667eea, #764ba2, #f093fb);
+        -webkit-background-clip: text; -webkit-text-fill-color: transparent;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+# ─── sidebar ─────────────────────────────────────────────────────────────
+st.sidebar.title("⚙️ Classification Controls")
+
+# Setup selector
+setup = st.sidebar.radio("Experiment Setup", ["Setup 1 (Natural)", "Setup 2 (Label Skew)"])
+setup_dir = "Setup_1" if "1" in setup else "Setup_2"
+models_dir = os.path.join(RESULTS, setup_dir, "models")
+
+# Check available models
+model_files = {}
+for name, fname in [("FedAvg", "fedavg_best.pth"),
+                     ("FedProx", "fedprox_best.pth"),
+                     ("QPSO-FL", "qpso_best.pth")]:
+    path = os.path.join(models_dir, fname)
+    if os.path.exists(path):
+        model_files[name] = path
+
+if not model_files:
+    st.error(f"No model weights found in `{models_dir}`. Please ensure .pth files are present.")
+    st.stop()
+
+selected_models = st.sidebar.multiselect(
+    "Compare Models",
+    list(model_files.keys()),
+    default=list(model_files.keys()),
+)
+
+st.sidebar.markdown("---")
+input_mode = st.sidebar.radio("Image Source", ["Upload Image", "Sample from Dataset"])
+
+# ─── title ───────────────────────────────────────────────────────────────
+st.markdown('<p class="hero-title">🧠 Brain Tumor Classification</p>',
+            unsafe_allow_html=True)
+st.markdown("**Federated Learning · SimpleCNN (~120K params) · "
+            "FedAvg vs FedProx vs QPSO-FL**")
+st.markdown("---")
+
+# ─── image input ─────────────────────────────────────────────────────────
+image = None
+
+if input_mode == "Upload Image":
+    uploaded = st.file_uploader(
+        "Upload a brain MRI image (JPG/PNG)",
+        type=["jpg", "jpeg", "png"],
+    )
+    if uploaded:
+        image = Image.open(uploaded).convert("RGB")
+
+elif input_mode == "Sample from Dataset":
+    # Try to find sample images from the results plots (confusion matrices show sample data)
+    # Or look for actual dataset images
+    sample_dirs = [
+        os.path.join(FL_ROOT, "data"),
+        os.path.join(FL_ROOT, "sample_images"),
+    ]
+
+    # Search for any sample images
+    sample_images = []
+    for d in sample_dirs:
+        if os.path.isdir(d):
+            for ext in ["*.jpg", "*.jpeg", "*.png"]:
+                sample_images.extend(glob.glob(os.path.join(d, "**", ext), recursive=True))
+
+    if sample_images:
+        # Group by class if possible
+        selected = st.selectbox("Select a sample image", sample_images,
+                                format_func=lambda x: os.path.basename(x))
+        image = Image.open(selected).convert("RGB")
+    else:
+        st.info(
+            "💡 No sample images found locally. You can either:\n"
+            "1. **Upload an image** using the sidebar option\n"
+            "2. **Add sample images** to `federated_learning/sample_images/` "
+            "(one subfolder per class: glioma/, meningioma/, notumor/, pituitary/)"
+        )
+
+# ─── inference ───────────────────────────────────────────────────────────
+if image is not None:
+    # Show the input image
+    st.subheader("📷 Input Image")
+    col_img, col_info = st.columns([1, 2])
+    with col_img:
+        st.image(image, caption="Input MRI", use_container_width=True)
+    with col_info:
+        w, h = image.size
+        st.markdown(f"**Resolution:** {w}×{h} → resized to {IMG_SIZE}×{IMG_SIZE}")
+        st.markdown(f"**Models:** {', '.join(selected_models)}")
+        st.markdown(f"**Setup:** {setup}")
+
+    st.markdown("---")
+    st.subheader("🔬 Classification Results")
+
+    if not selected_models:
+        st.warning("Select at least one model from the sidebar.")
+    else:
+        # Run all selected models
+        cols = st.columns(len(selected_models))
+        model_colors = {"FedAvg": "#1f77b4", "FedProx": "#ff7f0e", "QPSO-FL": "#2ca02c"}
+
+        results = {}
+        for idx, name in enumerate(selected_models):
+            with cols[idx]:
+                model, device = load_model(model_files[name])
+                pred_idx, probs = predict(model, device, image)
+                results[name] = (pred_idx, probs)
+                render_prediction_card(
+                    name,
+                    model_colors.get(name, "#666"),
+                    pred_idx, probs, image,
+                )
+
+        # Consensus section
+        if len(results) > 1:
+            st.markdown("---")
+            st.subheader("🤝 Model Consensus")
+
+            predictions = [CLASS_NAMES[r[0]] for r in results.values()]
+            unanimous = len(set(predictions)) == 1
+
+            if unanimous:
+                st.success(
+                    f"✅ **All {len(results)} models agree:** "
+                    f"{CLASS_ICONS[list(results.values())[0][0]]} "
+                    f"**{predictions[0]}**"
+                )
+            else:
+                # Majority vote
+                from collections import Counter
+                votes = Counter(predictions)
+                winner, count = votes.most_common(1)[0]
+                winner_idx = CLASS_NAMES.index(winner)
+                st.warning(
+                    f"⚠️ **Models disagree.** Majority vote ({count}/{len(results)}): "
+                    f"{CLASS_ICONS[winner_idx]} **{winner}**"
+                )
+
+                # Show disagreement details
+                for name, (pred_idx, probs) in results.items():
+                    emoji = "✅" if CLASS_NAMES[pred_idx] == winner else "❌"
+                    st.markdown(
+                        f"  {emoji} **{name}:** {CLASS_NAMES[pred_idx]} "
+                        f"({probs[pred_idx]*100:.1f}%)"
+                    )
+
+        # Average confidence across models
+        if len(results) > 1:
+            avg_probs = np.mean([r[1] for r in results.values()], axis=0)
+            ensemble_pred = int(np.argmax(avg_probs))
+
+            st.markdown("---")
+            st.subheader("📊 Ensemble Average (All Models)")
+            render_prediction_card(
+                "Ensemble Average",
+                "#E91E63",
+                ensemble_pred, avg_probs, image,
+            )
+
+else:
+    # Welcome state
+    st.markdown("""
+    ### How it works
+    1. **Choose a setup** — Natural heterogeneity (Setup 1) or Label Skew (Setup 2)
+    2. **Select models** — Compare FedAvg, FedProx, and QPSO-FL side by side
+    3. **Upload or select an image** — Any brain MRI (axial slice)
+    4. **See results** — Class prediction with confidence bars for each model
+
+    The models were trained using **federated learning** across 3 simulated hospitals,
+    each with different data distributions. The QPSO-FL model uses our novel
+    **Layer-by-Layer QPSO aggregation** for fairer global model performance.
+    """)
+
+    # Show model info cards
+    st.markdown("---")
+    info_cols = st.columns(3)
+    model_info = [
+        ("FedAvg", "#1f77b4", "Weighted average of client updates. Standard baseline."),
+        ("FedProx", "#ff7f0e", "Adds proximal regularization (μ=0.01) to prevent client drift."),
+        ("QPSO-FL", "#2ca02c", "Layer-by-layer quantum PSO with validation-loss fitness. Our contribution."),
+    ]
+    for col, (name, color, desc) in zip(info_cols, model_info):
+        with col:
+            st.markdown(
+                f'<div style="background:rgba(20,20,35,0.9);padding:20px;border-radius:12px;'
+                f'border-top:3px solid {color};">'
+                f'<h4 style="color:{color};margin:0 0 8px 0;">{name}</h4>'
+                f'<p style="color:#aaa;font-size:13px;margin:0;">{desc}</p></div>',
+                unsafe_allow_html=True,
+            )

pages/2_Slice_Viewer.py

@@ -0,0 +1,198 @@
+"""
+Slice-by-Slice Segmentation Viewer
+====================================
+View MRI slices with ground truth and AI prediction overlay.
+Supports all 4 modalities and 3 tumor sub-regions.
+"""
+
+import streamlit as st
+import os
+import sys
+import glob
+import nibabel as nib
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+
+# st.set_page_config(page_title="Slice Viewer", layout="wide")
+
+# ─── paths & inference ───────────────────────────────────────────────────
+APP_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if APP_DIR not in sys.path:
+    sys.path.insert(0, APP_DIR)
+from utils.inference import ensure_prediction, get_all_patients, DEMO_DIR
+
+
+def load_nifti(path):
+    if not os.path.exists(path):
+        return None
+    return nib.load(path).get_fdata()
+
+
+# ─── sidebar ─────────────────────────────────────────────────────────────
+st.sidebar.title("⚙️ Slice Viewer Controls")
+
+samples = get_all_patients()
+if not samples:
+    st.error("No demo data found. Please ensure patient volumes exist in demo_data/")
+    st.stop()
+
+selected_id = st.sidebar.selectbox("🧑‍⚕️ Patient", samples)
+
+modality = st.sidebar.selectbox(
+    "MRI Modality",
+    ["FLAIR", "T1", "T1ce", "T2"],
+    index=0,
+)
+MOD_MAP = {"T1": 0, "T1ce": 1, "T2": 2, "FLAIR": 3}
+
+overlay = st.sidebar.radio(
+    "Overlay",
+    ["AI Prediction", "Ground Truth", "Both (Side-by-Side)", "None"],
+    index=0,
+)
+
+overlay_alpha = st.sidebar.slider("Overlay Opacity", 0.1, 0.9, 0.5, 0.05)
+
+# ─── load data (run inference if needed) ─────────────────────────────────
+ensure_prediction(selected_id)
+
+img_path = os.path.join(DEMO_DIR, f"{selected_id}_image.nii.gz")
+pred_path = os.path.join(DEMO_DIR, f"{selected_id}_pred.nii.gz")
+lbl_path = os.path.join(DEMO_DIR, f"{selected_id}_label.nii.gz")
+
+img_data = load_nifti(img_path)    # (D, H, W, 4)
+pred_data = load_nifti(pred_path)  # (D, H, W, 3) channels: 0=TC, 1=WT, 2=ET
+lbl_data = load_nifti(lbl_path)    # (D, H, W) labels: 1=NCR, 2=ED, 4=ET
+
+if img_data is None:
+    st.error("Failed to load MRI volume.")
+    st.stop()
+
+depth = img_data.shape[0]
+
+# ─── title ───────────────────────────────────────────────────────────────
+st.title("🔬 Slice-by-Slice Segmentation Viewer")
+st.markdown(f"**Patient:** `{selected_id}` · **Modality:** {modality} · "
+            f"**Volume:** {img_data.shape[0]}×{img_data.shape[1]}×{img_data.shape[2]}")
+
+slice_idx = st.slider("Z-Axis Slice", 0, depth - 1, depth // 2)
+
+# ─── color maps ──────────────────────────────────────────────────────────
+# Tumor overlay: WT=green, TC=red, ET=yellow (matching 3D view)
+TUMOR_COLORS = np.array([
+    [0, 0, 0, 0],            # background (transparent)
+    [0.18, 0.80, 0.44, 1],   # WT - green
+    [0.91, 0.30, 0.24, 1],   # TC - red
+    [0.95, 0.77, 0.06, 1],   # ET - gold
+])
+
+
+def make_overlay_from_pred(pred_slice):
+    """Convert (H, W, 3) prediction channels to (H, W, 4) RGBA overlay."""
+    h, w = pred_slice.shape[:2]
+    overlay_img = np.zeros((h, w, 4))
+    # Order: WT first (background), then TC, then ET on top
+    wt = pred_slice[:, :, 1] > 0.5
+    tc = pred_slice[:, :, 0] > 0.5
+    et = pred_slice[:, :, 2] > 0.5
+    overlay_img[wt] = TUMOR_COLORS[1]
+    overlay_img[tc] = TUMOR_COLORS[2]
+    overlay_img[et] = TUMOR_COLORS[3]
+    return overlay_img
+
+
+def make_overlay_from_gt(lbl_slice):
+    """Convert integer label slice to (H, W, 4) RGBA overlay."""
+    h, w = lbl_slice.shape
+    overlay_img = np.zeros((h, w, 4))
+    wt = lbl_slice > 0
+    tc = (lbl_slice == 1) | (lbl_slice == 4)
+    et = lbl_slice == 4
+    overlay_img[wt] = TUMOR_COLORS[1]
+    overlay_img[tc] = TUMOR_COLORS[2]
+    overlay_img[et] = TUMOR_COLORS[3]
+    return overlay_img
+
+
+def render_slice(mri_slice, overlay_img, title, alpha):
+    """Render an MRI slice with optional overlay."""
+    fig, ax = plt.subplots(figsize=(6, 6))
+    ax.imshow(mri_slice, cmap="gray", origin="lower")
+    if overlay_img is not None:
+        ax.imshow(overlay_img, alpha=alpha, origin="lower")
+    ax.set_title(title, fontsize=14, color="white", fontweight="bold")
+    ax.axis("off")
+    fig.patch.set_facecolor("#0a0a14")
+    return fig
+
+
+# ─── render ──────────────────────────────────────────────────────────────
+mri_slice = img_data[slice_idx, :, :, MOD_MAP[modality]]
+
+if overlay == "None":
+    fig = render_slice(mri_slice, None, f"{modality} — Slice {slice_idx}", 0)
+    st.pyplot(fig)
+
+elif overlay == "AI Prediction":
+    if pred_data is not None:
+        ov = make_overlay_from_pred(pred_data[slice_idx])
+        fig = render_slice(mri_slice, ov, f"AI Prediction — Slice {slice_idx}",
+                           overlay_alpha)
+        st.pyplot(fig)
+    else:
+        st.warning("Prediction not available for this patient.")
+
+elif overlay == "Ground Truth":
+    if lbl_data is not None:
+        ov = make_overlay_from_gt(lbl_data[slice_idx])
+        fig = render_slice(mri_slice, ov, f"Ground Truth — Slice {slice_idx}",
+                           overlay_alpha)
+        st.pyplot(fig)
+    else:
+        st.warning("Ground truth not available.")
+
+elif overlay == "Both (Side-by-Side)":
+    col1, col2 = st.columns(2)
+    with col1:
+        if lbl_data is not None:
+            ov = make_overlay_from_gt(lbl_data[slice_idx])
+            fig = render_slice(mri_slice, ov, "Ground Truth", overlay_alpha)
+            st.pyplot(fig)
+        else:
+            st.info("No ground truth available.")
+    with col2:
+        if pred_data is not None:
+            ov = make_overlay_from_pred(pred_data[slice_idx])
+            fig = render_slice(mri_slice, ov, "AI Prediction", overlay_alpha)
+            st.pyplot(fig)
+        else:
+            st.info("No prediction available.")
+
+# ─── all modalities row ──────────────────────────────────────────────────
+with st.expander("📋 All 4 Modalities (this slice)", expanded=False):
+    cols = st.columns(4)
+    mod_names = ["T1", "T1ce", "T2", "FLAIR"]
+    for i, col in enumerate(cols):
+        with col:
+            fig, ax = plt.subplots(figsize=(3, 3))
+            ax.imshow(img_data[slice_idx, :, :, i], cmap="gray", origin="lower")
+            ax.set_title(mod_names[i], fontsize=11, color="white")
+            ax.axis("off")
+            fig.patch.set_facecolor("#0a0a14")
+            st.pyplot(fig)
+
+# ─── color legend ────────────────────────────────────────────────────────
+st.markdown("---")
+legend_cols = st.columns(3)
+labels = ["Whole Tumor (WT)", "Tumor Core (TC)", "Enhancing Tumor (ET)"]
+colors = ["#2ECC71", "#E74C3C", "#F1C40F"]
+for i, col in enumerate(legend_cols):
+    with col:
+        st.markdown(
+            f'<div style="display:flex;align-items:center;gap:8px;">'
+            f'<div style="width:16px;height:16px;background:{colors[i]};'
+            f'border-radius:3px;"></div>'
+            f'<span>{labels[i]}</span></div>',
+            unsafe_allow_html=True,
+        )

pages/3_3D_Visualization.py

@@ -0,0 +1,328 @@
+"""
+3D Interactive Brain Tumor Visualization
+=========================================
+Renders the brain + tumor regions as interactive 3D surfaces.
+Supports side-by-side Ground Truth vs AI Prediction comparison.
+Runs live inference if prediction doesn't exist yet.
+"""
+
+import streamlit as st
+import os
+import sys
+import glob
+import nibabel as nib
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+from skimage.measure import marching_cubes
+
+# st.set_page_config(page_title="3D Tumor Visualization", layout="wide")
+
+# ─── paths & inference ───────────────────────────────────────────────────
+APP_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if APP_DIR not in sys.path:
+    sys.path.insert(0, APP_DIR)
+from utils.inference import ensure_prediction, get_all_patients, DEMO_DIR
+
+
+def load_nifti(path):
+    if not os.path.exists(path):
+        return None
+    return nib.load(path).get_fdata()
+
+
+
+def extract_mesh(volume, level=0.5, step_size=2):
+    vol = volume[::step_size, ::step_size, ::step_size]
+    if vol.sum() == 0:
+        return None
+    try:
+        verts, faces, _, _ = marching_cubes(vol, level=level)
+        verts = verts * step_size
+        return verts, faces
+    except Exception:
+        return None
+
+
+def make_mesh_trace(volume, color, name, opacity, step_size, level=0.5,
+                    flatshading=True, scene="scene"):
+    result = extract_mesh(volume, level=level, step_size=step_size)
+    if result is None:
+        return None
+    verts, faces = result
+    x, y, z = verts.T
+    i, j, k = faces.T
+    return go.Mesh3d(
+        x=x, y=y, z=z, i=i, j=j, k=k,
+        color=color, opacity=opacity,
+        name=name, showlegend=True,
+        flatshading=flatshading,
+        lighting=dict(ambient=0.6, diffuse=0.7, specular=0.2, roughness=0.6),
+        lightposition=dict(x=100, y=200, z=300),
+        scene=scene,
+    )
+
+
+# ─── colors ──────────────────────────────────────────────────────────────
+PRED_COLORS = {
+    "Whole Tumor (WT)": "#2ECC71",   # emerald green
+    "Tumor Core (TC)":  "#E74C3C",   # vivid red
+    "Enhancing Tumor (ET)": "#F1C40F",  # bright gold
+}
+GT_COLORS = {
+    "Whole Tumor (WT)": "#1ABC9C",   # turquoise
+    "Tumor Core (TC)":  "#9B59B6",   # amethyst
+    "Enhancing Tumor (ET)": "#3498DB",  # ocean blue
+}
+PRED_CHANNELS = {"Whole Tumor (WT)": 1, "Tumor Core (TC)": 0, "Enhancing Tumor (ET)": 2}
+BRAIN_COLOR = "#D5D8DC"
+
+# ─── sidebar ─────────────────────────────────────────────────────────────
+st.sidebar.title("⚙️ 3D Controls")
+
+samples = get_all_patients()
+if not samples:
+    st.error("No processed prediction volumes found.")
+    st.stop()
+
+selected_id = st.sidebar.selectbox("🧑‍⚕️ Patient", samples)
+
+st.sidebar.markdown("---")
+view_mode = st.sidebar.radio(
+    "View Mode",
+    ["Prediction Only", "Ground Truth Only", "Side-by-Side Comparison"],
+    index=0,
+)
+
+st.sidebar.markdown("---")
+st.sidebar.subheader("🧠 Brain Surface")
+show_brain = st.sidebar.checkbox("Show Brain", value=True)
+brain_opacity = st.sidebar.slider("Brain Opacity", 0.02, 0.30, 0.08, 0.02)
+
+st.sidebar.subheader("🎯 Tumor")
+region_choice = st.sidebar.multiselect(
+    "Regions",
+    ["Whole Tumor (WT)", "Tumor Core (TC)", "Enhancing Tumor (ET)"],
+    default=["Whole Tumor (WT)", "Tumor Core (TC)", "Enhancing Tumor (ET)"],
+)
+tumor_opacity = st.sidebar.slider("Tumor Opacity", 0.20, 1.0, 0.70, 0.05)
+
+st.sidebar.markdown("---")
+step_size = st.sidebar.select_slider("Mesh Quality", options=[1, 2, 3, 4], value=2)
+
+# ─── load data (run inference if needed) ─────────────────────────────────
+ensure_prediction(selected_id)
+
+img_data = load_nifti(os.path.join(DEMO_DIR, f"{selected_id}_image.nii.gz"))
+pred_data = load_nifti(os.path.join(DEMO_DIR, f"{selected_id}_pred.nii.gz"))
+lbl_data = load_nifti(os.path.join(DEMO_DIR, f"{selected_id}_label.nii.gz"))
+
+# ─── title ───────────────────────────────────────────────────────────────
+st.title("🌐 3D Brain Tumor Visualization")
+st.markdown(f"**Patient:** `{selected_id}` · **Drag** to rotate · **Scroll** to zoom")
+
+# ─── color legend ────────────────────────────────────────────────────────
+legend_cols = st.columns(6)
+for idx, (region, color) in enumerate(PRED_COLORS.items()):
+    with legend_cols[idx]:
+        st.markdown(
+            f'<div style="display:flex;align-items:center;gap:6px;">'
+            f'<div style="width:14px;height:14px;background:{color};'
+            f'border-radius:3px;"></div><span style="font-size:13px;">Pred: {region}</span></div>',
+            unsafe_allow_html=True,
+        )
+if view_mode in ["Ground Truth Only", "Side-by-Side Comparison"]:
+    for idx, (region, color) in enumerate(GT_COLORS.items()):
+        with legend_cols[idx + 3]:
+            st.markdown(
+                f'<div style="display:flex;align-items:center;gap:6px;">'
+                f'<div style="width:14px;height:14px;background:{color};'
+                f'border-radius:3px;"></div><span style="font-size:13px;">GT: {region}</span></div>',
+                unsafe_allow_html=True,
+            )
+
+
+# ─── helper: build brain trace ───────────────────────────────────────────
+def build_brain_trace(scene_name="scene"):
+    if img_data is None:
+        return None
+    flair = img_data[:, :, :, 3]
+    flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)
+    brain_mask = (flair_norm > 0.15).astype(float)
+    return make_mesh_trace(
+        brain_mask, BRAIN_COLOR, "Brain",
+        opacity=brain_opacity,
+        step_size=max(step_size, 2),
+        flatshading=False, scene=scene_name,
+    )
+
+
+def build_pred_traces(scene_name="scene"):
+    traces = []
+    if pred_data is None:
+        return traces
+    for region in region_choice:
+        ch = PRED_CHANNELS[region]
+        vol = pred_data[:, :, :, ch]
+        t = make_mesh_trace(vol, PRED_COLORS[region], f"Pred: {region}",
+                            tumor_opacity, step_size, scene=scene_name)
+        if t:
+            traces.append(t)
+    return traces
+
+
+def build_gt_traces(scene_name="scene"):
+    traces = []
+    if lbl_data is None:
+        return traces
+    gt_masks = {
+        "Whole Tumor (WT)": (lbl_data > 0).astype(float),
+        "Tumor Core (TC)": ((lbl_data == 1) | (lbl_data == 4)).astype(float),
+        "Enhancing Tumor (ET)": (lbl_data == 4).astype(float),
+    }
+    for region in region_choice:
+        vol = gt_masks[region]
+        t = make_mesh_trace(vol, GT_COLORS[region], f"GT: {region}",
+                            tumor_opacity, step_size, scene=scene_name)
+        if t:
+            traces.append(t)
+    return traces
+
+
+SCENE_LAYOUT = dict(
+    xaxis=dict(visible=False),
+    yaxis=dict(visible=False),
+    zaxis=dict(visible=False),
+    bgcolor="rgb(10, 10, 20)",
+    aspectmode="data",
+    camera=dict(eye=dict(x=1.6, y=1.0, z=0.8), up=dict(x=0, y=0, z=1)),
+)
+
+# ─── render ──────────────────────────────────────────────────────────────
+
+if view_mode == "Side-by-Side Comparison":
+    # Two 3D plots: GT on the left, Prediction on the right
+    col_left, col_right = st.columns(2)
+
+    with col_left:
+        st.markdown("### 🟢 Ground Truth")
+        gt_traces = []
+        if show_brain:
+            bt = build_brain_trace("scene")
+            if bt:
+                gt_traces.append(bt)
+        gt_traces.extend(build_gt_traces("scene"))
+
+        if gt_traces:
+            fig_gt = go.Figure(data=gt_traces)
+            fig_gt.update_layout(
+                scene=SCENE_LAYOUT,
+                margin=dict(l=0, r=0, t=0, b=0),
+                height=600,
+                paper_bgcolor="rgb(10, 10, 20)",
+                legend=dict(font=dict(color="white", size=11),
+                            bgcolor="rgba(20,20,40,0.8)", x=0.01, y=0.99),
+            )
+            st.plotly_chart(fig_gt, width="stretch")
+        else:
+            st.info("No ground truth data available for this patient.")
+
+    with col_right:
+        st.markdown("### 🔴 AI Prediction")
+        pred_traces = []
+        if show_brain:
+            bt = build_brain_trace("scene")
+            if bt:
+                pred_traces.append(bt)
+        pred_traces.extend(build_pred_traces("scene"))
+
+        if pred_traces:
+            fig_pred = go.Figure(data=pred_traces)
+            fig_pred.update_layout(
+                scene=SCENE_LAYOUT,
+                margin=dict(l=0, r=0, t=0, b=0),
+                height=600,
+                paper_bgcolor="rgb(10, 10, 20)",
+                legend=dict(font=dict(color="white", size=11),
+                            bgcolor="rgba(20,20,40,0.8)", x=0.01, y=0.99),
+            )
+            st.plotly_chart(fig_pred, width="stretch")
+        else:
+            st.warning("No prediction data available.")
+
+else:
+    # Single 3D view
+    all_traces = []
+    if show_brain:
+        bt = build_brain_trace("scene")
+        if bt:
+            all_traces.append(bt)
+
+    if view_mode == "Prediction Only":
+        all_traces.extend(build_pred_traces("scene"))
+    elif view_mode == "Ground Truth Only":
+        all_traces.extend(build_gt_traces("scene"))
+
+    if not all_traces:
+        st.warning("Nothing to render. Check that data exists and regions are selected.")
+        st.stop()
+
+    fig = go.Figure(data=all_traces)
+    fig.update_layout(
+        scene=SCENE_LAYOUT,
+        margin=dict(l=0, r=0, t=0, b=0),
+        height=750,
+        paper_bgcolor="rgb(10, 10, 20)",
+        legend=dict(font=dict(color="white", size=13),
+                    bgcolor="rgba(20,20,40,0.85)",
+                    bordercolor="rgba(100,100,140,0.5)", borderwidth=1,
+                    x=0.01, y=0.99),
+    )
+    st.plotly_chart(fig, width="stretch")
+
+# ─── volume stats ────────────────────────────────────────────────────────
+st.markdown("---")
+st.subheader("📊 Tumor Volume Statistics")
+
+if pred_data is not None:
+    cols = st.columns(3)
+    for idx, (region, ch) in enumerate(PRED_CHANNELS.items()):
+        vol = pred_data[:, :, :, ch]
+        voxel_count = int(vol.sum())
+        volume_cc = voxel_count / 1000.0
+        color = PRED_COLORS[region]
+        with cols[idx]:
+            st.markdown(
+                f'<div style="background:rgba(30,30,50,0.8);padding:14px;'
+                f'border-radius:10px;border-left:4px solid {color};">'
+                f'<div style="color:{color};font-size:13px;font-weight:600;">{region}</div>'
+                f'<div style="color:white;font-size:26px;font-weight:700;">{volume_cc:.1f} cm³</div>'
+                f'<div style="color:#888;font-size:11px;">{voxel_count:,} voxels</div></div>',
+                unsafe_allow_html=True,
+            )
+
+# ─── dice ────────────────────────────────────────────────────────────────
+if lbl_data is not None and pred_data is not None:
+    st.markdown("---")
+    st.subheader("🔬 Dice Scores")
+    cols = st.columns(3)
+    gt_masks_dice = {
+        "Whole Tumor (WT)": (lbl_data > 0).astype(float),
+        "Tumor Core (TC)": ((lbl_data == 1) | (lbl_data == 4)).astype(float),
+        "Enhancing Tumor (ET)": (lbl_data == 4).astype(float),
+    }
+    for idx, (region, ch) in enumerate(PRED_CHANNELS.items()):
+        p = pred_data[:, :, :, ch]
+        g = gt_masks_dice[region]
+        dice = (2.0 * (p * g).sum()) / (p.sum() + g.sum() + 1e-8)
+        color = PRED_COLORS[region]
+        grade = "Excellent" if dice > 0.8 else "Good" if dice > 0.6 else "Fair"
+        with cols[idx]:
+            st.markdown(
+                f'<div style="background:rgba(30,30,50,0.8);padding:14px;'
+                f'border-radius:10px;border-left:4px solid {color};">'
+                f'<div style="color:{color};font-size:13px;">{region}</div>'
+                f'<div style="color:white;font-size:30px;font-weight:700;">{dice:.4f}</div>'
+                f'<div style="color:#888;">{grade}</div></div>',
+                unsafe_allow_html=True,
+            )

requirements.txt

@@ -0,0 +1,9 @@
+torch>=1.12.0
+torchvision>=0.13.0
+numpy>=1.21.0
+Pillow>=9.0.0
+plotly>=5.10.0
+streamlit>=1.20.0
+matplotlib>=3.5.0
+nibabel>=4.0.0
+monai>=1.0.0

results/Setup_1/models/fedavg_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d48f93f0f3ed5e4074a24f6bc645bd55d56068467573902be606e94a75363f0
+size 631956

results/Setup_1/models/fedprox_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7b6c27da43d10bbee7275de1d34eb417f3c7867694e8aabdf9ee3deebebbc52
+size 631987

results/Setup_1/models/qpso_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bcd1b0bcaf795d4519a8457dba2e778a25adbda949adf337e35bd6b185f38ef
+size 631894

results/Setup_2/models/fedavg_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d48f93f0f3ed5e4074a24f6bc645bd55d56068467573902be606e94a75363f0
+size 631956

results/Setup_2/models/fedprox_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7b6c27da43d10bbee7275de1d34eb417f3c7867694e8aabdf9ee3deebebbc52
+size 631987

results/Setup_2/models/qpso_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bcd1b0bcaf795d4519a8457dba2e778a25adbda949adf337e35bd6b185f38ef
+size 631894

src/streamlit_app.py

@@ -0,0 +1,40 @@
+import altair as alt
+import numpy as np
+import pandas as pd
+import streamlit as st
+
+"""
+# Welcome to Streamlit!
+
+Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
+If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
+forums](https://discuss.streamlit.io).
+
+In the meantime, below is an example of what you can do with just a few lines of code:
+"""
+
+num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
+num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+
+indices = np.linspace(0, 1, num_points)
+theta = 2 * np.pi * num_turns * indices
+radius = indices
+
+x = radius * np.cos(theta)
+y = radius * np.sin(theta)
+
+df = pd.DataFrame({
+    "x": x,
+    "y": y,
+    "idx": indices,
+    "rand": np.random.randn(num_points),
+})
+
+st.altair_chart(alt.Chart(df, height=700, width=700)
+    .mark_point(filled=True)
+    .encode(
+        x=alt.X("x", axis=None),
+        y=alt.Y("y", axis=None),
+        color=alt.Color("idx", legend=None, scale=alt.Scale()),
+        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
+    ))

utils/inference.py

@@ -0,0 +1,212 @@
+"""
+Shared inference module for 3D brain tumor segmentation.
+Loads the AttentionUnet model and runs sliding_window_inference
+on patients that don't have pre-computed predictions.
+"""
+
+import os
+import shutil
+import numpy as np
+import nibabel as nib
+import streamlit as st
+import torch
+from monai.inferers import sliding_window_inference
+from monai.networks.nets import AttentionUnet
+from monai.transforms import (
+    Compose,
+    LoadImaged,
+    NormalizeIntensityd,
+    Orientationd,
+    Spacingd,
+    EnsureChannelFirstd,
+    EnsureTyped,
+)
+
+
+# ─── paths ───────────────────────────────────────────────────────────────
+SEG_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "segmentation"))
+DEMO_DIR = os.path.join(SEG_DIR, "demo_data")
+# streamlit_app/ is inside segmentation/, so go up one level to reach segmentation/
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+
+# Model checkpoint — prefer refined model (better calibration)
+# 1. Check streamlit_app/ (where refined model lives)
+# 2. Check segmentation/ (parent dir, where base model lives)
+_THIS_DIR = os.path.dirname(os.path.dirname(__file__))
+_candidates = [
+    os.path.join(_THIS_DIR, "best_metric_model_refined.pth"),          # streamlit_app/
+    os.path.join(PROJECT_ROOT, "best_metric_model_refined.pth"),       # segmentation/
+    os.path.join(_THIS_DIR, "best_metric_model.pth"),                  # streamlit_app/
+    os.path.join(PROJECT_ROOT, "best_metric_model.pth"),               # segmentation/
+]
+CKPT_PATH = None
+for _c in _candidates:
+    if os.path.exists(_c):
+        CKPT_PATH = _c
+        break
+
+# MONAI transforms — must match training exactly
+INFERENCE_TRANSFORMS = Compose([
+    LoadImaged(keys=["image", "label"]),
+    EnsureChannelFirstd(keys=["image", "label"]),
+    EnsureTyped(keys=["image", "label"]),
+    Orientationd(keys=["image", "label"], axcodes="RAS"),
+    Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
+    NormalizeIntensityd(keys=["image"], nonzero=True, channel_wise=True),
+])
+
+# Transforms for image-only (no label available)
+INFERENCE_TRANSFORMS_IMG_ONLY = Compose([
+    LoadImaged(keys=["image"]),
+    EnsureChannelFirstd(keys=["image"]),
+    EnsureTyped(keys=["image"]),
+    Orientationd(keys=["image"], axcodes="RAS"),
+    Spacingd(keys=["image"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear",)),
+    NormalizeIntensityd(keys=["image"], nonzero=True, channel_wise=True),
+])
+
+
+@st.cache_resource
+def load_seg_model():
+    """Load the 3D Attention U-Net model (cached across sessions)."""
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model = AttentionUnet(
+        spatial_dims=3,
+        in_channels=4,
+        out_channels=3,
+        channels=(16, 32, 64, 128, 256),
+        strides=(2, 2, 2, 2),
+    ).to(device)
+
+    if os.path.exists(CKPT_PATH):
+        try:
+            model.load_state_dict(torch.load(CKPT_PATH, map_location=device))
+            model.eval()
+            return model, device
+        except Exception as e:
+            st.error(f"Failed to load model weights: {e}")
+            return None, None
+    else:
+        st.error(f"Model checkpoint not found at {CKPT_PATH}")
+        return None, None
+
+
+def ensure_prediction(patient_id):
+    """
+    Ensure that the prediction volume exists for a patient.
+    If _pred.nii.gz already exists, returns True immediately.
+    Otherwise, runs live inference using the exact same MONAI
+    transforms as the training pipeline.
+    """
+    pred_path = os.path.join(DEMO_DIR, f"{patient_id}_pred.nii.gz")
+    img_path = os.path.join(DEMO_DIR, f"{patient_id}_image.nii.gz")
+    lbl_path = os.path.join(DEMO_DIR, f"{patient_id}_label.nii.gz")
+
+    # Already have prediction — skip
+    if os.path.exists(pred_path) and os.path.exists(img_path):
+        return True
+
+    # Check if raw MRI modalities exist in patient subfolder
+    p_dir = os.path.join(DEMO_DIR, patient_id)
+    if not os.path.isdir(p_dir):
+        return False
+
+    # Build file paths (same order as extract_demo_data.py: t1, t1ce, t2, flair)
+    mod_paths = {
+        "t1":    os.path.join(p_dir, f"{patient_id}_t1.nii.gz"),
+        "t1ce":  os.path.join(p_dir, f"{patient_id}_t1ce.nii.gz"),
+        "t2":    os.path.join(p_dir, f"{patient_id}_t2.nii.gz"),
+        "flair": os.path.join(p_dir, f"{patient_id}_flair.nii.gz"),
+    }
+    seg_path = os.path.join(p_dir, f"{patient_id}_seg.nii.gz")
+
+    for m, mp in mod_paths.items():
+        if not os.path.exists(mp):
+            st.warning(f"Missing modality: {m} at {mp}")
+            return False
+
+    # ─── Run live inference ──────────────────────────────────────────
+    st.info(f"🧠 **Running AI Inference** on `{patient_id}`... This may take 30-60 seconds.")
+    progress = st.progress(0)
+    status = st.empty()
+
+    try:
+        # Build MONAI data dict (image is a list of 4 modality paths)
+        has_label = os.path.exists(seg_path)
+        data_dict = {
+            "image": [mod_paths["t1"], mod_paths["t1ce"], mod_paths["t2"], mod_paths["flair"]],
+        }
+        if has_label:
+            data_dict["label"] = seg_path
+
+        # Apply MONAI transforms (Orientation, Spacing, Normalize — matching training)
+        status.text("Loading & preprocessing with MONAI transforms...")
+        if has_label:
+            sample_data = INFERENCE_TRANSFORMS(data_dict)
+        else:
+            sample_data = INFERENCE_TRANSFORMS_IMG_ONLY(data_dict)
+        progress.progress(30)
+
+        # Run model inference
+        status.text("Running 3D U-Net inference (sliding window)...")
+        model, device = load_seg_model()
+        if model is None:
+            return False
+
+        inputs = sample_data["image"].unsqueeze(0).to(device)  # (1, 4, D, H, W)
+        with torch.no_grad():
+            outputs = sliding_window_inference(inputs, (96, 96, 96), 4, model)
+            outputs = (outputs.sigmoid() > 0.5).float()
+        progress.progress(80)
+
+        # Save processed image volume (D, H, W, 4)
+        status.text("Saving results...")
+        img_np = inputs[0].cpu().numpy().transpose(1, 2, 3, 0)
+        nib.save(nib.Nifti1Image(img_np, affine=np.eye(4)), img_path)
+
+        # Save prediction (D, H, W, 3)
+        pred_np = outputs[0].cpu().numpy().transpose(1, 2, 3, 0)
+        nib.save(nib.Nifti1Image(pred_np, affine=np.eye(4)), pred_path)
+
+        # Save ground truth label (D, H, W)
+        if has_label:
+            lbl_np = sample_data["label"][0].cpu().numpy()
+            nib.save(nib.Nifti1Image(lbl_np.astype(np.float32), affine=np.eye(4)), lbl_path)
+        elif not os.path.exists(lbl_path):
+            empty = np.zeros(pred_np.shape[:3])
+            nib.save(nib.Nifti1Image(empty.astype(np.float32), affine=np.eye(4)), lbl_path)
+
+        progress.progress(100)
+        status.text("✅ Inference complete!")
+        return True
+
+    except Exception as e:
+        st.error(f"Inference failed: {e}")
+        import traceback
+        st.code(traceback.format_exc())
+        return False
+
+
+def get_all_patients():
+    """
+    Return all patient IDs that have either pre-computed predictions
+    OR raw MRI data (can be inferred on-demand).
+    """
+    patients = set()
+
+    # Patients with pre-computed predictions
+    import glob
+    for f in glob.glob(os.path.join(DEMO_DIR, "*_pred.nii.gz")):
+        pid = os.path.basename(f).replace("_pred.nii.gz", "")
+        patients.add(pid)
+
+    # Patients with raw MRI data (subfolder with modality files)
+    if os.path.isdir(DEMO_DIR):
+        for d in os.listdir(DEMO_DIR):
+            full = os.path.join(DEMO_DIR, d)
+            if os.path.isdir(full) and d.startswith("BraTS"):
+                # Check it has at least the flair file
+                if os.path.exists(os.path.join(full, f"{d}_flair.nii.gz")):
+                    patients.add(d)
+
+    return sorted(patients)