Update app.py

app.py

@@ -1,6 +1,4 @@
 from pathlib import Path
-import json
-import re
 import traceback
 
 import gradio as gr
@@ -9,7 +7,6 @@ import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
 
-
 from model_adapter import (
     MODEL_SPECS,
     load_roi_model,
@@ -19,9 +16,11 @@ from model_adapter import (
 )
 from roi_extractor import extract_rois_from_t1
 
+
 APP_DIR = Path(__file__).resolve().parent
 MODELS_DIR = APP_DIR / "models"
 
+
 DISPLAY_CLASSES = [
     "No Impairment",
     "Very Mild Impairment",
@@ -37,16 +36,19 @@ CLASS_COLORS = {
     "Moderate Impairment": "#ef4444",
 }
 
+
 TASK_OPTIONS = [
     "Progression",
     "Development",
 ]
 
+
 MRI_STRATEGY_OPTIONS = [
     "Multi-Region ROI",
     "Hippocampal ROI",
 ]
 
+
 CORE_FEATURES = [
     "age_at_visit",
     "sex",
@@ -65,20 +67,29 @@ CORE_FEATURES = [
     "entorhinal_thickness_l",
     "entorhinal_thickness_r",
 ]
+
+
 DEFAULT_CORE_ROWS = [[name, ""] for name in CORE_FEATURES]
 
 
 PROJECT_MD = r"""
-# AD-Stage-Net v2: Multimodal Alzheimer’s Disease Prediction
+# AD-Stage-Net: Multimodal Alzheimer’s Disease Prediction
 
-**AD-Stage-Net Multimodal** extends the original MRI classification project by combining
-anatomy-focused structural MRI with clinical and cognitive information.
+**AD-Stage-Net Multimodal** combines structural MRI with clinical and cognitive
+information to evaluate future Alzheimer’s disease outcomes.
 
-Upload a **whole-volume T1 MRI** (`.nii` / `.nii.gz`) to run either the hippocampal
-or multi-region multimodal model. The application automatically performs anatomical
-segmentation and ROI extraction before multimodal inference.
+**Original Staging Model:** <https://huggingface.co/spaces/katelynhur/AD-Stage-Net>
+
+Upload a **whole-volume T1 MRI** (`.nii` / `.nii.gz`). The application automatically
+performs anatomical segmentation and ROI extraction before multimodal inference.
+
+### Prediction Tasks
+
+- **Progression** — predicts future disease progression
+- **Development** — predicts future development of Alzheimer’s-related impairment
+
+### MRI Strategies
 
-**New multimodal MRI choices**
 - **Multi-Region ROI:** hippocampus + entorhinal cortex + amygdala + ventricles
 - **Hippocampal ROI:** focused hippocampal representation
 
@@ -86,31 +97,43 @@ segmentation and ROI extraction before multimodal inference.
 """
 
 
-# ---------------------------------------------------------------------
-# Shared display helpers
-# ---------------------------------------------------------------------
+# ============================================================
+# Basic helpers
+# ============================================================
+
 def _to_float(value):
     if value is None:
         return np.nan
+
     if isinstance(value, str):
         s = value.strip()
+
         if not s:
             return np.nan
+
         low = s.lower()
+
         if low in {"m", "male"}:
             return 1.0
+
         if low in {"f", "female"}:
             return 0.0
+
         if low in {"yes", "true", "positive"}:
             return 1.0
+
         if low in {"no", "false", "negative"}:
             return 0.0
+
         try:
             return float(s)
+
         except ValueError:
             return np.nan
+
     try:
         return float(value)
+
     except Exception:
         return np.nan
 
@@ -118,109 +141,203 @@ def _to_float(value):
 def _core_table_to_dict(core_df):
     if core_df is None:
         return {}
+
     if not isinstance(core_df, pd.DataFrame):
-        core_df = pd.DataFrame(core_df, columns=["Feature", "Value"])
-    out = {}
+        core_df = pd.DataFrame(
+            core_df,
+            columns=["Feature", "Value"],
+        )
+
+    output = {}
+
     for _, row in core_df.iterrows():
         key = str(row.iloc[0]).strip()
+
         if key:
-            out[key] = _to_float(row.iloc[1])
-    return out
+            output[key] = _to_float(row.iloc[1])
+
+    return output
 
 
 def _read_extra_features(extra_csv):
     if extra_csv is None:
         return {}
+
     df = pd.read_csv(Path(extra_csv))
+
     if df.empty:
         return {}
 
-    lower_cols = {str(c).strip().lower(): c for c in df.columns}
+    lower_cols = {
+        str(c).strip().lower(): c
+        for c in df.columns
+    }
+
     if "feature" in lower_cols and "value" in lower_cols:
-        fk = lower_cols["feature"]
-        vk = lower_cols["value"]
+        feature_col = lower_cols["feature"]
+        value_col = lower_cols["value"]
+
         return {
-            str(r[fk]).strip(): _to_float(r[vk])
-            for _, r in df.iterrows()
-            if str(r[fk]).strip()
+            str(row[feature_col]).strip(): _to_float(row[value_col])
+            for _, row in df.iterrows()
+            if str(row[feature_col]).strip()
         }
 
     row = df.iloc[0]
-    return {str(c): _to_float(row[c]) for c in df.columns}
-
-
-def _make_prob_chart(probs):
-    vals = [float(probs.get(c, 0.0)) * 100.0 for c in DISPLAY_CLASSES]
-    fig = go.Figure(go.Bar(
-        x=vals,
-        y=DISPLAY_CLASSES,
-        orientation="h",
-        marker_color=[CLASS_COLORS[c] for c in DISPLAY_CLASSES],
-        text=[f"{v:.1f}%" for v in vals],
-        textposition="outside",
-    ))
+
+    return {
+        str(column): _to_float(row[column])
+        for column in df.columns
+    }
+
+
+def _is_nifti(path):
+    name = str(path).lower()
+
+    return (
+        name.endswith(".nii")
+        or name.endswith(".nii.gz")
+    )
+
+
+# ============================================================
+# Result display
+# ============================================================
+
+def _make_prob_chart(probabilities):
+    values = [
+        float(probabilities.get(stage, 0.0)) * 100
+        for stage in DISPLAY_CLASSES
+    ]
+
+    fig = go.Figure(
+        go.Bar(
+            x=values,
+            y=DISPLAY_CLASSES,
+            orientation="h",
+            marker_color=[
+                CLASS_COLORS[stage]
+                for stage in DISPLAY_CLASSES
+            ],
+            text=[
+                f"{value:.1f}%"
+                for value in values
+            ],
+            textposition="outside",
+        )
+    )
+
     fig.update_layout(
         xaxis_title="Probability (%)",
         xaxis_range=[0, 105],
-        height=270,
-        margin=dict(l=15, r=20, t=15, b=25),
+        height=280,
+        margin=dict(
+            l=20,
+            r=30,
+            t=20,
+            b=30,
+        ),
         plot_bgcolor="rgba(0,0,0,0)",
         paper_bgcolor="rgba(0,0,0,0)",
     )
+
     return fig
 
 
-def _stage_html(stage, confidence, model_name):
-    color = CLASS_COLORS.get(stage, "#64748b")
+def _stage_html(
+    stage,
+    confidence,
+    task_name,
+    strategy_name,
+):
+    color = CLASS_COLORS.get(
+        stage,
+        "#64748b",
+    )
+
     return f"""
-    <div style="border:1px solid {color};border-radius:14px;padding:16px 18px;
-                background:{color}12;text-align:center;">
-      <div style="font-size:.78rem;opacity:.7;">Predicted Stage</div>
-      <div style="font-size:1.55rem;font-weight:750;color:{color};margin-top:3px;">{stage}</div>
-      <div style="font-size:.95rem;margin-top:4px;">{confidence*100:.1f}% confidence</div>
-      <div style="font-size:.74rem;opacity:.62;margin-top:6px;">{model_name}</div>
+    <div style="
+        border:1px solid {color};
+        border-radius:14px;
+        padding:18px 20px;
+        background:{color}12;
+        text-align:center;
+        min-height:120px;
+    ">
+
+        <div style="
+            font-size:.8rem;
+            opacity:.7;
+        ">
+            Predicted Result
+        </div>
+
+        <div style="
+            font-size:1.6rem;
+            font-weight:750;
+            color:{color};
+            margin-top:4px;
+        ">
+            {stage}
+        </div>
+
+        <div style="
+            font-size:1rem;
+            margin-top:5px;
+        ">
+            {confidence * 100:.1f}% confidence
+        </div>
+
+        <div style="
+            font-size:.78rem;
+            opacity:.65;
+            margin-top:7px;
+        ">
+            {task_name} · {strategy_name}
+        </div>
+
     </div>
     """
 
 
 def _softmax_or_normalize(raw):
-    arr = np.asarray(raw, dtype=np.float64).reshape(-1)
-    if len(arr) != 4:
-        raise ValueError(f"Expected four outputs, got {np.asarray(raw).shape}.")
-    if np.all(arr >= 0) and 0.98 <= arr.sum() <= 1.02:
-        return arr / arr.sum()
-    e = np.exp(arr - arr.max())
-    return e / e.sum()
+    arr = np.asarray(
+        raw,
+        dtype=np.float64,
+    ).reshape(-1)
 
+    if len(arr) != 4:
+        raise ValueError(
+            f"Expected four outputs, got {np.asarray(raw).shape}."
+        )
 
+    if (
+        np.all(arr >= 0)
+        and 0.98 <= arr.sum() <= 1.02
+    ):
+        return arr / arr.sum()
 
-def _is_nifti(path):
-    name = str(path).lower()
-    return name.endswith(".nii") or name.endswith(".nii.gz")
+    exp = np.exp(
+        arr - arr.max()
+    )
 
+    return exp / exp.sum()
 
 
-# ---------------------------------------------------------------------
-# Task / MRI-strategy routing
-# ---------------------------------------------------------------------
-def _resolve_model_key(task_choice, mri_strategy):
-    """
-    Resolve the selected prediction task and MRI representation.
+# ============================================================
+# Task + MRI strategy routing
+# ============================================================
 
-    Current Stage uses the existing deployment keys:
-      - multi_region
-      - hippocampal
+def _resolve_model_key(
+    task_choice,
+    mri_strategy,
+):
 
-    Progression and Development expect corresponding entries in
-    model_config.json. This keeps the UI honest: the option is visible, but the
-    app will not silently reuse the staging model for a different task.
-    """
-    strategy_key = (
-        "multi_region"
-        if mri_strategy.startswith("Multi-Region")
-        else "hippocampal"
-    )
+    if mri_strategy == "Multi-Region ROI":
+        strategy_key = "multi_region"
 
+    else:
+        strategy_key = "hippocampal"
 
     if task_choice == "Progression":
         return f"progression_{strategy_key}"
@@ -228,7 +345,9 @@ def _resolve_model_key(task_choice, mri_strategy):
     if task_choice == "Development":
         return f"development_{strategy_key}"
 
-    raise gr.Error(f"Unknown prediction task: {task_choice}")
+    raise gr.Error(
+        f"Unknown prediction task: {task_choice}"
+    )
 
 
 def _task_output_description(task_choice):
@@ -247,304 +366,730 @@ def _task_output_description(task_choice):
     return ""
 
 
-# ---------------------------------------------------------------------
-# Multimodal inference
-# ---------------------------------------------------------------------
+# ============================================================
+# Model helpers
+# ============================================================
+
 def _required_files_message(model_key):
     spec = MODEL_SPECS[model_key]
+
     missing = []
-    for filename in (spec["roi_checkpoint"], spec["fusion_model"]):
-        if not (MODELS_DIR / filename).exists():
+
+    for filename in (
+        spec["roi_checkpoint"],
+        spec["fusion_model"],
+    ):
+
+        if not (
+            MODELS_DIR / filename
+        ).exists():
+
             missing.append(filename)
+
     if not missing:
         return ""
-    return "Missing model file(s): " + ", ".join(f"models/{x}" for x in missing)
 
+    return (
+        "Missing model file(s): "
+        + ", ".join(
+            f"models/{filename}"
+            for filename in missing
+        )
+    )
+
+
+def _build_feature_frame(
+    booster,
+    clinical_values,
+    embedding,
+    spec,
+):
 
-def _build_feature_frame(booster, clinical_values, embedding, spec):
     feature_names = booster.feature_name()
-    values = {name: np.nan for name in feature_names}
 
-    for name, val in clinical_values.items():
+    values = {
+        name: np.nan
+        for name in feature_names
+    }
+
+    for name, value in clinical_values.items():
         if name in values:
-            values[name] = val
+            values[name] = value
+
+    configured_names = spec.get(
+        "embedding_feature_names",
+        [],
+    )
 
-    configured = spec.get("embedding_feature_names", [])
-    if configured:
-        if len(configured) != len(embedding):
+    if configured_names:
+
+        if len(configured_names) != len(embedding):
             raise ValueError(
-                f"Configured embedding list has {len(configured)} features, "
-                f"but the encoder returned {len(embedding)} values."
+                "Configured MRI embedding contains "
+                f"{len(configured_names)} names, but encoder returned "
+                f"{len(embedding)} values."
             )
-        for name, val in zip(configured, embedding):
+
+        for name, value in zip(
+            configured_names,
+            embedding,
+        ):
+
             if name in values:
-                values[name] = float(val)
+                values[name] = float(value)
+
     else:
-        prefixes = ("img_", "mri_", "emb_", "embedding_", "roi_")
-        candidates = [n for n in feature_names if n.lower().startswith(prefixes)]
+
+        prefixes = (
+            "img_",
+            "mri_",
+            "emb_",
+            "embedding_",
+            "roi_",
+        )
+
+        candidates = [
+            name
+            for name in feature_names
+            if name.lower().startswith(prefixes)
+        ]
+
         if len(candidates) != len(embedding):
             raise ValueError(
-                "Could not map the MRI embedding into the LightGBM model. "
-                "Fill embedding_feature_names in model_config.json with the exact "
-                "training feature names."
+                "Could not automatically map the MRI embedding "
+                "into the LightGBM model. Add the exact MRI "
+                "embedding feature names to embedding_feature_names "
+                "in model_config.json."
             )
-        for name, val in zip(candidates, embedding):
-            values[name] = float(val)
 
-    return pd.DataFrame([[values[n] for n in feature_names]], columns=feature_names)
+        for name, value in zip(
+            candidates,
+            embedding,
+        ):
+            values[name] = float(value)
+
+    return pd.DataFrame(
+        [
+            [
+                values[name]
+                for name in feature_names
+            ]
+        ],
+        columns=feature_names,
+    )
+
+
+# ============================================================
+# Multimodal inference
+# ============================================================
 
+def _run_multimodal(
+    task_choice,
+    mri_strategy,
+    mri_path,
+    core_df,
+    extra_csv,
+):
 
-def _run_multimodal(task_choice, mri_strategy, mri_path, core_df, extra_csv):
     if not _is_nifti(mri_path):
         raise gr.Error(
-            "The hippocampal and multi-region models require a whole-volume "
-            "T1 MRI in .nii or .nii.gz format."
+            "Please upload a whole-volume T1 MRI "
+            "in .nii or .nii.gz format."
         )
 
-    model_key = _resolve_model_key(task_choice, mri_strategy)
+    model_key = _resolve_model_key(
+        task_choice,
+        mri_strategy,
+    )
 
     if model_key not in MODEL_SPECS:
         raise gr.Error(
-            f"No deployment model is configured yet for '{task_choice}' with "
-            f"'{mri_strategy}'. Add a '{model_key}' entry to model_config.json "
-            "using the exact trained checkpoint and fusion model for that task. "
-            "The app will not substitute the staging model."
+            f"No deployment model is configured for "
+            f"'{task_choice}' using '{mri_strategy}'. "
+            f"Add '{model_key}' to model_config.json "
+            "using the exact trained model files."
         )
 
     spec = MODEL_SPECS[model_key]
 
-    missing = _required_files_message(model_key)
+    missing = _required_files_message(
+        model_key
+    )
+
     if missing:
         raise gr.Error(missing)
 
+    # --------------------------------------------------------
+    # Automatic segmentation + ROI extraction
+    # --------------------------------------------------------
+
     roi_arrays, roi_qc = extract_rois_from_t1(
         mri_path,
-        crop_shape=tuple(spec["crop_shape"]),
+        crop_shape=tuple(
+            spec["crop_shape"]
+        ),
         device="cuda",
     )
 
-    selected = {
-        region: normalize_roi(
+    selected_rois = {}
+
+    for region in spec["regions"]:
+
+        selected_rois[region] = normalize_roi(
             roi_arrays[region],
-            expected_shape=tuple(spec["crop_shape"]),
+            expected_shape=tuple(
+                spec["crop_shape"]
+            ),
         )
-        for region in spec["regions"]
-    }
 
-    encoder = load_roi_model(MODELS_DIR / spec["roi_checkpoint"])
-    embedding = extract_embedding(encoder, selected, spec)
-    booster = load_fusion_model(MODELS_DIR / spec["fusion_model"])
+    # --------------------------------------------------------
+    # MRI encoder
+    # --------------------------------------------------------
+
+    encoder = load_roi_model(
+        MODELS_DIR
+        / spec["roi_checkpoint"]
+    )
+
+    embedding = extract_embedding(
+        encoder,
+        selected_rois,
+        spec,
+    )
+
+    # --------------------------------------------------------
+    # Clinical / cognitive fusion
+    # --------------------------------------------------------
+
+    booster = load_fusion_model(
+        MODELS_DIR
+        / spec["fusion_model"]
+    )
+
+    clinical_values = _core_table_to_dict(
+        core_df
+    )
+
+    clinical_values.update(
+        _read_extra_features(
+            extra_csv
+        )
+    )
+
+    X = _build_feature_frame(
+        booster,
+        clinical_values,
+        embedding,
+        spec,
+    )
+
+    raw_prediction = booster.predict(X)
+
+    raw_array = np.asarray(
+        raw_prediction
+    )
+
+    if raw_array.ndim == 2:
+        raw_array = raw_array[0]
+
+    probabilities = _softmax_or_normalize(
+        raw_array
+    )
+
+    probability_map = dict(
+        zip(
+            DISPLAY_CLASSES,
+            probabilities.tolist(),
+        )
+    )
+
+    top_index = int(
+        np.argmax(probabilities)
+    )
+
+    stage = DISPLAY_CLASSES[
+        top_index
+    ]
 
-    clinical = _core_table_to_dict(core_df)
-    clinical.update(_read_extra_features(extra_csv))
+    confidence = float(
+        probabilities[top_index]
+    )
 
-    X = _build_feature_frame(booster, clinical, embedding, spec)
-    raw = booster.predict(X)
-    raw_arr = np.asarray(raw)
-    if raw_arr.ndim == 2:
-        raw_arr = raw_arr[0]
-    probs = _softmax_or_normalize(raw_arr)
+    regions = ", ".join(
+        spec["regions"]
+    )
 
-    top_i = int(np.argmax(probs))
-    stage = DISPLAY_CLASSES[top_i]
-    confidence = float(probs[top_i])
-    prob_map = dict(zip(DISPLAY_CLASSES, probs.tolist()))
+    filled_features = int(
+        np.sum(
+            ~pd.isna(X.iloc[0])
+        )
+    )
 
-    regions = ", ".join(spec["regions"])
-    filled = int(np.sum(~pd.isna(X.iloc[0])))
-    total = len(X.columns)
+    total_features = len(
+        X.columns
+    )
 
-    qc = "; ".join(
-        f"{r}: {roi_qc[r]['label_voxels']} voxels"
-        for r in spec["regions"]
+    qc_text = "; ".join(
+        (
+            f"{region}: "
+            f"{roi_qc[region]['label_voxels']} voxels"
+        )
+        for region in spec["regions"]
     )
 
     status = (
         f"**Prediction task:** {task_choice}  \n"
         f"**Task meaning:** {_task_output_description(task_choice)}  \n"
+        f"**MRI strategy:** {mri_strategy}  \n"
         f"**Input:** whole-volume T1 MRI  \n"
-        f"**Automatic processing:** FastSurfer → 64×64×64 ROI extraction  \n"
+        f"**Automatic processing:** FastSurfer → ROI extraction  \n"
         f"**Regions used:** {regions}  \n"
-        f"**Fusion features populated:** {filled}/{total}  \n"
-        f"**ROI QC:** {qc}"
+        f"**ROI size:** "
+        f"{spec['crop_shape'][0]}×"
+        f"{spec['crop_shape'][1]}×"
+        f"{spec['crop_shape'][2]}  \n"
+        f"**Fusion features populated:** "
+        f"{filled_features}/{total_features}  \n"
+        f"**ROI QC:** {qc_text}"
+    )
+
+    return (
+        stage,
+        confidence,
+        probability_map,
+        status,
+        X,
     )
-    return stage, confidence, prob_map, status, X
 
 
-# ---------------------------------------------------------------------
-# Main ZeroGPU function
-# ---------------------------------------------------------------------
+# ============================================================
+# ZeroGPU wrapper
+# ============================================================
+
 @spaces.GPU(duration=300)
-def run_inference(task_choice, mri_strategy, mri_file, core_df, extra_csv):
+def run_inference(
+    task_choice,
+    mri_strategy,
+    mri_file,
+    core_df,
+    extra_csv,
+):
+
     if mri_file is None:
-        raise gr.Error("Please upload an MRI.")
+        raise gr.Error(
+            "Please upload a whole-volume T1 MRI."
+        )
 
     try:
-        stage, confidence, probs, status, X = _run_multimodal(
-            task_choice, mri_strategy, mri_file, core_df, extra_csv
+
+        (
+            stage,
+            confidence,
+            probabilities,
+            status,
+            feature_frame,
+        ) = _run_multimodal(
+            task_choice,
+            mri_strategy,
+            mri_file,
+            core_df,
+            extra_csv,
         )
 
         return (
-            _stage_html(stage, confidence, f"{task_choice} · {mri_strategy}"),
-            _make_prob_chart(probs),
+            _stage_html(
+                stage,
+                confidence,
+                task_choice,
+                mri_strategy,
+            ),
+            _make_prob_chart(
+                probabilities
+            ),
             status,
-            X,
+            feature_frame,
         )
+
     except gr.Error:
         raise
-    except Exception as exc:
+
+    except Exception as error:
+
         traceback.print_exc()
-        raise gr.Error(f"Inference failed: {exc}")
 
+        raise gr.Error(
+            f"Inference failed: {error}"
+        )
+
+
+# ============================================================
+# Dynamic UI helper
+# ============================================================
+
+def selection_help(
+    task_choice,
+    mri_strategy,
+):
+
+    if mri_strategy == "Hippocampal ROI":
+
+        region_text = (
+            "**Automatic ROI:** hippocampus"
+        )
 
-def selection_help(task_choice, mri_strategy):
-    if mri_strategy.startswith("Hippocampal"):
-        region_text = "**Automatic ROI:** hippocampus"
     else:
+
         region_text = (
-            "**Automatic ROIs:** hippocampus + entorhinal cortex + "
-            "amygdala + ventricles"
+            "**Automatic ROIs:** hippocampus + "
+            "entorhinal cortex + amygdala + ventricles"
         )
 
     if task_choice == "Progression":
-        task_text = "**Task:** predict future disease progression"
+
+        task_text = (
+            "**Task:** predict future disease progression"
+        )
+
     else:
+
         task_text = (
-            "**Task:** predict future development of Alzheimer’s-related impairment"
+            "**Task:** predict future development "
+            "of Alzheimer’s-related impairment"
         )
 
     return (
         f"{task_text}  \n"
         "**Input:** whole-volume T1 MRI (`.nii` / `.nii.gz`)  \n"
         f"{region_text}  \n"
-        "**Clinical/cognitive data:** used in fusion"
+        "**Clinical/cognitive data:** used in multimodal fusion"
     )
 
 
+# ============================================================
+# Styling
+# ============================================================
+
 CSS = """
 .gradio-container {
-    max-width: 1260px !important;
+    max-width: 1600px !important;
+    width: 95vw !important;
+    margin: 0 auto !important;
 }
+
 .hero-card {
     background: white;
     border: 1px solid #e5e7eb;
     border-radius: 14px;
-    padding: 22px 26px;
-    margin-bottom: 8px;
+    padding: 24px 28px;
+    margin-bottom: 12px;
+}
+
+.control-card {
+    border: 1px solid #e5e7eb;
+    border-radius: 12px;
+    padding: 18px;
+    background: white;
 }
-.input-card, .performance-card {
+
+.info-card {
     border: 1px solid #e5e7eb;
     border-radius: 12px;
-    padding: 12px;
+    padding: 18px 20px;
+    background: white;
 }
+
 #result-card {
-    min-height: 120px;
+    min-height: 130px;
+}
+
+footer {
+    visibility: hidden;
 }
-footer {visibility: hidden;}
 """
 
 
-with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="AD-Stage-Net Multimodal") as demo:
-    # Same top hero-card structure as the original app.
+# ============================================================
+# App layout
+# ============================================================
+
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css=CSS,
+    title="AD-Stage-Net Multimodal",
+) as demo:
+
+    # ========================================================
+    # Header
+    # ========================================================
+
     with gr.Row():
+
         with gr.Column():
-            gr.Markdown(PROJECT_MD, elem_classes=["hero-card"])
 
-    # Same basic 3:2 main-page layout as the original app.
-    with gr.Row():
-        with gr.Column(scale=3):
-            with gr.Row():
-                with gr.Column():
-                    task_choice = gr.Dropdown(
-                        choices=TASK_OPTIONS,
-                        value="Progression",
-                        label="Prediction Task",
-                        allow_custom_value=False,
-                    )
+            gr.Markdown(
+                PROJECT_MD,
+                elem_classes=[
+                    "hero-card"
+                ],
+            )
 
-                    mri_strategy = gr.Dropdown(
-                        choices=MRI_STRATEGY_OPTIONS,
-                        value="Multi-Region ROI",
-                        label="MRI Strategy",
-                        allow_custom_value=False,
-                    )
+    # ========================================================
+    # Main layout
+    # ========================================================
 
-                    model_note = gr.Markdown(
-                        selection_help("Progression", "Multi-Region ROI")
-                    )
+    with gr.Row(
+        equal_height=False
+    ):
 
-                    extra_csv = gr.File(
-                        label="Optional cognitive-feature CSV",
-                        file_types=[".csv"],
-                        type="filepath",
-                    )
+        # ====================================================
+        # LEFT SIDE — inference controls
+        # ====================================================
 
-                with gr.Column():
-                    mri_file = gr.File(
-                        label="Upload whole-volume T1 MRI",
-                        file_types=[".nii", ".gz"],
-                        type="filepath",
-                    )
-                    gr.Markdown(
-                        "**Supported:** NIfTI whole-volume MRI (`.nii`, `.nii.gz`)."
+        with gr.Column(
+            scale=4,
+            min_width=650,
+        ):
+
+            # ------------------------------------------------
+            # Prediction controls
+            # ------------------------------------------------
+
+            gr.Markdown(
+                "## Run a Prediction"
+            )
+
+            with gr.Group(
+                elem_classes=[
+                    "control-card"
+                ]
+            ):
+
+                task_choice = gr.Dropdown(
+                    choices=TASK_OPTIONS,
+                    value="Progression",
+                    label="Prediction Task",
+                    allow_custom_value=False,
+                )
+
+                mri_strategy = gr.Dropdown(
+                    choices=MRI_STRATEGY_OPTIONS,
+                    value="Multi-Region ROI",
+                    label="MRI Strategy",
+                    allow_custom_value=False,
+                )
+
+                model_note = gr.Markdown(
+                    selection_help(
+                        "Progression",
+                        "Multi-Region ROI",
                     )
-                    clear_btn = gr.Button("Clear file")
+                )
+
+            # ------------------------------------------------
+            # MRI upload
+            # ------------------------------------------------
+
+            gr.Markdown(
+                "### 1. Upload MRI"
+            )
+
+            mri_file = gr.File(
+                label="Upload whole-volume T1 MRI",
+                file_types=[
+                    ".nii",
+                    ".gz",
+                ],
+                type="filepath",
+                height=180,
+            )
 
-            gr.Markdown("### Clinical / Cognitive Information")
             gr.Markdown(
-                "Used by the hippocampal and multi-region multimodal models. "
-                "Leave unavailable measurements blank."
+                "**Supported:** NIfTI whole-volume MRI "
+                "(`.nii`, `.nii.gz`)."
+            )
+
+            clear_btn = gr.Button(
+                "Clear MRI",
+                variant="secondary",
             )
+
+            # ------------------------------------------------
+            # Clinical / cognitive information
+            # ------------------------------------------------
+
+            gr.Markdown(
+                "### 2. Clinical / Cognitive Information"
+            )
+
+            gr.Markdown(
+                "Enter the measurements available for the participant. "
+                "Leave unavailable values blank."
+            )
+
             core_df = gr.Dataframe(
-                headers=["Feature", "Value"],
-                datatype=["str", "str"],
+                headers=[
+                    "Feature",
+                    "Value",
+                ],
+                datatype=[
+                    "str",
+                    "str",
+                ],
                 value=DEFAULT_CORE_ROWS,
-                row_count=(len(DEFAULT_CORE_ROWS), "fixed"),
-                col_count=(2, "fixed"),
+                row_count=(
+                    len(DEFAULT_CORE_ROWS),
+                    "fixed",
+                ),
+                col_count=(
+                    2,
+                    "fixed",
+                ),
                 interactive=True,
                 label="Clinical / cognitive features",
+                height=420,
+            )
+
+            extra_csv = gr.File(
+                label="Optional complete cognitive-feature CSV",
+                file_types=[
+                    ".csv"
+                ],
+                type="filepath",
             )
 
-            run_btn = gr.Button("Run inference", variant="primary", size="lg")
+            # ------------------------------------------------
+            # Run
+            # ------------------------------------------------
+
+            gr.Markdown(
+                "### 3. Run Model"
+            )
+
+            run_btn = gr.Button(
+                "Run inference",
+                variant="primary",
+                size="lg",
+            )
+
+        # ====================================================
+        # RIGHT SIDE — information
+        # ====================================================
+
+        with gr.Column(
+            scale=2,
+            min_width=350,
+        ):
+
+            with gr.Group(
+                elem_classes=[
+                    "info-card"
+                ]
+            ):
 
-            gr.Markdown("### Prediction")
-            with gr.Row():
-                stage_output = gr.HTML(elem_id="result-card")
-                chart_output = gr.Plot()
-            status_output = gr.Markdown()
+                gr.Markdown(
+                    """
+## How the pipeline works
 
-            with gr.Accordion("Show exact multimodal fusion input", open=False):
-                feature_frame = gr.Dataframe(interactive=False)
+**Whole-volume T1 MRI**
 
-        with gr.Column(scale=2):
-            gr.Markdown("""
-### How the new pipeline works
+↓  
 
-**Whole-volume T1 MRI**  
-→ automatic FastSurfer segmentation  
-→ automatic anatomical ROI extraction  
-→ MRI embedding  
-→ clinical/cognitive fusion  
-→ prediction
+**FastSurfer segmentation**
 
-The **multi-region** representation uses the hippocampus, entorhinal cortex,
-amygdala, and ventricles.
+↓
 
-The **hippocampal** representation uses the hippocampal ROI only.
+**Automatic ROI extraction**
 
-Clinical and cognitive information is combined with the MRI representation
-for the selected prediction task.
-""")
+↓
+
+**MRI embedding**
+
+↓
+
+**Clinical + cognitive fusion**
+
+↓
+
+**Prediction**
+
+---
+
+### Prediction Tasks
+
+**Progression**  
+Predicts future disease progression using information available at the current visit.
+
+**Development**  
+Predicts whether the participant will develop future Alzheimer’s-related impairment.
+
+---
+
+### MRI Strategies
+
+**Multi-Region ROI**  
+Uses the hippocampus, entorhinal cortex, amygdala, and ventricles.
+
+**Hippocampal ROI**  
+Uses the hippocampal region only.
+"""
+                )
+
+    # ========================================================
+    # Prediction results — full width
+    # ========================================================
+
+    gr.Markdown("---")
+
+    gr.Markdown(
+        "## Prediction Results"
+    )
+
+    with gr.Row():
+
+        with gr.Column(
+            scale=2
+        ):
+
+            stage_output = gr.HTML(
+                elem_id="result-card"
+            )
+
+        with gr.Column(
+            scale=3
+        ):
+
+            chart_output = gr.Plot()
+
+    status_output = gr.Markdown()
+
+    with gr.Accordion(
+        "Show exact multimodal fusion input",
+        open=False,
+    ):
+
+        feature_frame = gr.Dataframe(
+            interactive=False
+        )
+
+    # ========================================================
+    # Events
+    # ========================================================
 
     task_choice.change(
         fn=selection_help,
-        inputs=[task_choice, mri_strategy],
+        inputs=[
+            task_choice,
+            mri_strategy,
+        ],
         outputs=model_note,
     )
 
     mri_strategy.change(
         fn=selection_help,
-        inputs=[task_choice, mri_strategy],
+        inputs=[
+            task_choice,
+            mri_strategy,
+        ],
         outputs=model_note,
     )
 
@@ -556,10 +1101,25 @@ for the selected prediction task.
 
     run_btn.click(
         fn=run_inference,
-        inputs=[task_choice, mri_strategy, mri_file, core_df, extra_csv],
-        outputs=[stage_output, chart_output, status_output, feature_frame],
+        inputs=[
+            task_choice,
+            mri_strategy,
+            mri_file,
+            core_df,
+            extra_csv,
+        ],
+        outputs=[
+            stage_output,
+            chart_output,
+            status_output,
+            feature_frame,
+        ],
     )
 
 
+# ============================================================
+# Launch
+# ============================================================
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()