Update app.py

app.py

@@ -1,25 +1,26 @@
 import os
 import pickle
 import warnings
+import tempfile
 import numpy as np
 import torch
 import torch.nn as nn
 from math import sqrt
 import gradio as gr
+import nibabel as nib
+from sklearn.preprocessing import MinMaxScaler
 
 # ══════════════════════════════════════════════════════════════════════════════
-# 1. Model definitions (must exactly match training code)
+# 1. Model definitions
 # ══════════════════════════════════════════════════════════════════════════════
 
 class Sine(nn.Module):
-    def __init__(self, w0: float = 1.0):
+    def __init__(self, w0=1.0):
         super().__init__()
         self.w0 = w0
-
     def forward(self, x):
         return torch.sin(self.w0 * x)
 
-
 class SirenLayer(nn.Module):
     def __init__(self, dim_in, dim_out, w0=30.0, c=6.0,
                  is_first=False, use_bias=True, activation=None):
@@ -30,11 +31,9 @@ class SirenLayer(nn.Module):
         if use_bias:
             nn.init.uniform_(self.linear.bias, -w_std, w_std)
         self.activation = Sine(w0) if activation is None else activation
-
     def forward(self, x):
         return self.activation(self.linear(x))
 
-
 class Siren(nn.Module):
     def __init__(self, dim_in, dim_hidden, dim_out, num_layers,
                  w0=30.0, w0_initial=30.0, use_bias=True, final_activation=None):
@@ -50,149 +49,313 @@ class Siren(nn.Module):
         act = nn.Identity() if final_activation is None else final_activation
         self.last_layer = SirenLayer(dim_hidden, dim_out, w0=w0,
                                      use_bias=use_bias, activation=act)
-
     def forward(self, x):
         return self.last_layer(self.net(x))
 
-
 class SirenMRIModel(nn.Module):
-    """One Siren per axial slice, bundled into a single nn.Module."""
-
     def __init__(self, config):
         super().__init__()
         self.config = config
         self.models = nn.ModuleList([
-            Siren(
-                dim_in=2,
-                dim_hidden=config["layer_size"],
-                dim_out=config["vols"],
-                num_layers=config["num_layers"],
-                w0=config["w0"],
-                w0_initial=config["w0_initial"],
-            )
+            Siren(dim_in=2, dim_hidden=config["layer_size"],
+                  dim_out=config["vols"], num_layers=config["num_layers"],
+                  w0=config["w0"], w0_initial=config["w0_initial"])
             for _ in range(config["sz"])
         ])
-
     def forward(self, coords, slice_idx):
         return self.models[slice_idx](coords)
 
-
 # ══════════════════════════════════════════════════════════════════════════════
-# 2. Load model + scalers at startup
+# 2. Load pretrained model
 # ══════════════════════════════════════════════════════════════════════════════
 
 def load_assets():
-    model_path   = "sirenMRI_full_model_final.pt"
-    scalers_path = "scalers.pkl"
-
-    for path in (model_path, scalers_path):
-        if not os.path.exists(path):
-            raise FileNotFoundError(
-                f"Required file not found: {path}\n"
-                "Upload sirenMRI_full_model_final.pt and scalers.pkl "
-                "to the root of your Space."
-            )
-
-    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
-    config     = checkpoint["config"]
-
-    model = SirenMRIModel(config)
-    model.load_state_dict(checkpoint["model_state_dict"])
-    model.eval()
-
+    model_path, scalers_path = "sirenMRI_full_model_final.pt", "scalers.pkl"
+    for p in (model_path, scalers_path):
+        if not os.path.exists(p):
+            raise FileNotFoundError(f"Missing: {p}")
+    ckpt    = torch.load(model_path, map_location="cpu", weights_only=False)
+    cfg     = ckpt["config"]
+    mdl     = SirenMRIModel(cfg)
+    mdl.load_state_dict(ckpt["model_state_dict"])
+    mdl.eval()
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         with open(scalers_path, "rb") as f:
-            scalers = pickle.load(f)
+            scl = pickle.load(f)
+    return mdl, scl, cfg, ckpt["input_shape"]
 
-    return model, scalers, config, checkpoint["input_shape"]
-
-
-print("Loading model…")
+print("⏳ Loading model…")
 model, scalers, config, input_shape = load_assets()
 sx, sy, sz, vols = input_shape
-print(f"Model loaded — shape: {sx}x{sy}x{sz}, volumes: {vols}")
-
+print(f"✅ Model ready — {sx}×{sy}×{sz}, {vols} volumes")
 
 # ══════════════���═══════════════════════════════════════════════════════════════
-# 3. Inference
+# 3. Helper: normalise to uint8
 # ══════════════════════════════════════════════════════════════════════════════
 
-def reconstruct(slice_idx: int, vol_idx: int):
-    """Reconstruct one axial slice and return it as a displayable image."""
-    slice_idx = int(slice_idx)
-    vol_idx   = int(vol_idx)
+def to_uint8(arr):
+    a, b = arr.min(), arr.max()
+    return ((arr - a) / (b - a + 1e-8) * 255).astype(np.uint8)
 
-    xs = torch.linspace(-1, 1, sx)
-    ys = torch.linspace(-1, 1, sy)
-    grid_x, grid_y = torch.meshgrid(xs, ys, indexing="ij")
-    coords = torch.stack([grid_x.reshape(-1), grid_y.reshape(-1)], dim=-1)
+def to_coords(h, w):
+    xs = torch.linspace(-1, 1, h)
+    ys = torch.linspace(-1, 1, w)
+    gx, gy = torch.meshgrid(xs, ys, indexing="ij")
+    return torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=-1)
 
-    with torch.no_grad():
-        pred = model(coords, slice_idx).numpy()      # (sx*sy, vols)
+# ══════════════════════════════════════════════════════════════════════════════
+# 4a. Reconstruct from pretrained model
+# ══════════════════════════════════════════════════════════════════════════════
 
-    # Inverse MinMaxScaler — pure numpy, no sklearn call
+def reconstruct_pretrained(slice_idx, vol_idx):
+    slice_idx, vol_idx = int(slice_idx), int(vol_idx)
+    coords = to_coords(sx, sy)
+    with torch.no_grad():
+        pred = model(coords, slice_idx).numpy()
     scaler   = scalers[slice_idx]
     data_min = np.array(scaler.data_min_, dtype=np.float32)
     data_max = np.array(scaler.data_max_, dtype=np.float32)
-    pred     = pred * (data_max - data_min) + data_min  # (sx*sy, vols)
-
-    img = pred.reshape(sx, sy, vols)[:, :, vol_idx]     # (sx, sy)
-
-    # Normalise to [0, 255] uint8 for display
-    img_min, img_max = img.min(), img.max()
-    img_norm = (img - img_min) / (img_max - img_min + 1e-8)
-    img_uint8 = (img_norm * 255).astype(np.uint8)
-
-    info = (
-        f"Slice {slice_idx} | Volume {vol_idx} | "
-        f"Shape: {img.shape} | "
-        f"Intensity range: [{img_min:.3f}, {img_max:.3f}]"
+    pred     = pred * (data_max - data_min) + data_min
+    recon    = pred.reshape(sx, sy, vols)[:, :, vol_idx]
+    img_min, img_max = recon.min(), recon.max()
+    stats = (
+        f"📐 Shape: {recon.shape}  |  "
+        f"📊 Intensity: [{img_min:.3f}, {img_max:.3f}]  |  "
+        f"🧠 Slice {slice_idx}  |  📡 Volume {vol_idx}"
     )
-    return img_uint8, info
-
+    return to_uint8(recon), stats
 
 # ══════════════════════════════════════════════════════════════════════════════
-# 4. Gradio UI
+# 4b. Compress & reconstruct user-uploaded NIfTI
 # ══════════════════════════════════════════════════════════════════════════════
 
-with gr.Blocks(title="SIREN MRI Compression") as demo:
-    gr.Markdown(
-        """
-        # 🧠 Physics-Informed SIREN MRI Compression
-        Reconstruct diffusion MRI slices from the compressed SIREN neural representation.
-        Select an axial slice and diffusion volume, then click **Reconstruct**.
-        """
-    )
+def compress_and_compare(nifti_file, slice_idx, vol_idx, num_iters, lr):
+    if nifti_file is None:
+        return None, None, "⚠️ Please upload a NIfTI file first."
 
-    with gr.Row():
-        slice_slider = gr.Slider(
-            minimum=0, maximum=sz - 1, step=1, value=sz // 2,
-            label=f"Axial Slice (0 – {sz - 1})"
-        )
-        vol_slider = gr.Slider(
-            minimum=0, maximum=vols - 1, step=1, value=0,
-            label=f"Diffusion Volume / b-value index (0 – {vols - 1})"
-        )
-
-    reconstruct_btn = gr.Button("▶ Reconstruct", variant="primary")
-
-    with gr.Row():
-        output_image = gr.Image(label="Reconstructed Slice", type="numpy")
-        output_info  = gr.Textbox(label="Info", lines=2, interactive=False)
-
-    reconstruct_btn.click(
-        fn=reconstruct,
-        inputs=[slice_slider, vol_slider],
-        outputs=[output_image, output_info],
+    slice_idx = int(slice_idx)
+    vol_idx   = int(vol_idx)
+    num_iters = int(num_iters)
+
+    try:
+        nii      = nib.load(nifti_file.name)
+        img_data = nii.get_fdata().astype(np.float32)
+    except Exception as e:
+        return None, None, f"❌ Failed to load NIfTI: {e}"
+
+    # Handle 3D (single volume) or 4D
+    if img_data.ndim == 3:
+        img_data = img_data[..., np.newaxis]
+    if img_data.ndim != 4:
+        return None, None, "❌ Expected a 3D or 4D NIfTI file."
+
+    ux, uy, uz, uvols = img_data.shape
+    slice_idx = min(slice_idx, uz - 1)
+    vol_idx   = min(vol_idx,   uvols - 1)
+
+    # ── Original slice ────────────────────────────────────────────────────────
+    orig_slice = img_data[:, :, slice_idx, vol_idx]
+    orig_img   = to_uint8(orig_slice)
+
+    # ── Quick SIREN fit on this one slice ─────────────────────────────────────
+    img_slice = np.transpose(img_data[:, :, slice_idx, :], (2, 0, 1))  # (vols, h, w)
+    features  = img_slice.reshape(uvols, -1).T                          # (h*w, vols)
+
+    scaler_u  = MinMaxScaler(feature_range=(0, 1))
+    features_scaled = scaler_u.fit_transform(features).astype(np.float32)
+
+    siren_u = Siren(dim_in=2, dim_hidden=config["layer_size"],
+                    dim_out=uvols, num_layers=config["num_layers"],
+                    w0=config["w0"], w0_initial=config["w0_initial"])
+    opt = torch.optim.Adam(siren_u.parameters(), lr=float(lr))
+    loss_fn = nn.MSELoss()
+
+    coords_u  = to_coords(ux, uy)
+    feat_t    = torch.from_numpy(features_scaled)
+
+    siren_u.train()
+    losses = []
+    for it in range(num_iters):
+        opt.zero_grad()
+        pred = siren_u(coords_u)
+        loss = loss_fn(pred, feat_t)
+        loss.backward()
+        opt.step()
+        losses.append(loss.item())
+
+    # ── Reconstruct ───────────────────────────────────────────────────────────
+    siren_u.eval()
+    with torch.no_grad():
+        pred_np = siren_u(coords_u).numpy()
+
+    pred_inv    = scaler_u.inverse_transform(pred_np)  # (h*w, vols)
+    recon_slice = pred_inv.reshape(ux, uy, uvols)[:, :, vol_idx]
+    recon_img   = to_uint8(recon_slice)
+
+    # ── PSNR ──────────────────────────────────────────────────────────────────
+    mse   = np.mean((orig_slice - recon_slice) ** 2)
+    o_max = orig_slice.max()
+    psnr  = 20 * np.log10(o_max / (np.sqrt(mse) + 1e-8)) if o_max > 0 else float("nan")
+    final_loss = losses[-1] if losses else float("nan")
+
+    stats = (
+        f"📐 Image: {ux}×{uy}×{uz}, {uvols} volumes  |  "
+        f"🎯 Slice {slice_idx}, Volume {vol_idx}\n"
+        f"📉 Final loss: {final_loss:.6f}  |  "
+        f"📡 PSNR: {psnr:.2f} dB  |  "
+        f"🔁 Iterations: {num_iters}"
     )
+    return orig_img, recon_img, stats
 
-    gr.Markdown(
-        f"**Model config** — Type: `{config['model'].upper()}` | "
-        f"Layers: `{config['num_layers']}` | "
-        f"Hidden size: `{config['layer_size']}` | "
-        f"w0: `{config['w0']}` | "
-        f"Slices: `{sz}` | Volumes: `{vols}`"
-    )
+# ══════════════════════════════════════════════════════════════════════════════
+# 5. Gradio UI
+# ══════════════════════════════════════════════════════════════════════════════
+
+CSS = """
+:root { --primary: #6366f1; --bg: #0f0f1a; --card: #1a1a2e; --border: #2d2d4e; }
+body, .gradio-container { background: var(--bg) !important; color: #e2e8f0 !important; }
+.gr-button-primary { background: linear-gradient(135deg,#6366f1,#8b5cf6) !important;
+    border: none !important; border-radius: 10px !important; font-weight: 700 !important;
+    letter-spacing: .5px; transition: transform .15s, box-shadow .15s; }
+.gr-button-primary:hover { transform: translateY(-2px);
+    box-shadow: 0 8px 25px rgba(99,102,241,.45) !important; }
+.gr-panel, .gr-box { background: var(--card) !important;
+    border: 1px solid var(--border) !important; border-radius: 14px !important; }
+.gr-input, .gr-slider { background: #12122a !important; border-color: var(--border) !important; }
+label { color: #a5b4fc !important; font-weight: 600 !important; }
+.gr-markdown h1 { background: linear-gradient(135deg,#6366f1,#a78bfa);
+    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
+    font-size: 2.2rem !important; font-weight: 800 !important; }
+.gr-markdown h2 { color: #a5b4fc !important; font-size: 1.1rem !important; }
+.tab-nav button { color: #a5b4fc !important; border-radius: 8px 8px 0 0 !important; }
+.tab-nav button.selected { background: var(--card) !important;
+    border-bottom: 2px solid #6366f1 !important; color: #fff !important; }
+footer { display: none !important; }
+"""
+
+with gr.Blocks(css=CSS, title="SIREN MRI Compression") as demo:
+
+    # ── Header ────────────────────────────────────────────────────────────────
+    gr.Markdown("""
+# 🧠 Physics-Informed SIREN MRI Compression
+## Neural implicit representation for diffusion MRI — compress, reconstruct & compare
+---
+""")
+
+    with gr.Tabs():
+
+        # ══════════════════════════════════════════════════════════════════════
+        # TAB 1 — Pretrained model explorer
+        # ══════════════════════════════════════════════════════════════════════
+        with gr.Tab("🔬 Explore Pretrained Model"):
+            gr.Markdown("""
+### Explore the model trained on the MGH-1010 diffusion dataset
+Adjust the sliders and click **Reconstruct** to visualise any slice and volume.
+""")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    sl1 = gr.Slider(0, sz-1,  value=sz//2,  step=1, label=f"Axial Slice  (0 – {sz-1})")
+                    vl1 = gr.Slider(0, vols-1, value=0,     step=1, label=f"Diffusion Volume  (0 – {vols-1})")
+                    btn1 = gr.Button("▶  Reconstruct", variant="primary")
+                    stats1 = gr.Textbox(label="Statistics", lines=2, interactive=False)
+
+                    gr.Markdown(f"""
+---
+**Model config**
+| Parameter | Value |
+|---|---|
+| Type | `{config['model'].upper()}` |
+| Layers | `{config['num_layers']}` |
+| Hidden size | `{config['layer_size']}` |
+| w₀ | `{config['w0']}` |
+| Slices | `{sz}` |
+| Volumes | `{vols}` |
+""")
+
+                with gr.Column(scale=2):
+                    out1 = gr.Image(label="Reconstructed Slice", type="numpy",
+                                    elem_id="recon_img", height=420)
+
+            btn1.click(reconstruct_pretrained,
+                       inputs=[sl1, vl1],
+                       outputs=[out1, stats1])
+
+        # ══════════════════════════════════════════════════════════════════════
+        # TAB 2 — Upload your own NIfTI
+        # ══════════════════════════════════════════════════════════════════════
+        with gr.Tab("📂 Upload & Compress Your Own MRI"):
+            gr.Markdown("""
+### Upload your own diffusion MRI in NIfTI format
+The app will fit a SIREN network to the selected slice on-the-fly and show you
+**original vs reconstructed** side by side.  
+> ⚠️ For speed, only the selected slice is fitted. Use more iterations for better quality.
+""")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    nifti_upload = gr.File(
+                        label="Upload NIfTI file (.nii or .nii.gz)",
+                        file_types=[".nii", ".gz"],
+                    )
+                    sl2 = gr.Slider(0, 200, value=50,  step=1,  label="Axial Slice")
+                    vl2 = gr.Slider(0, 551, value=0,   step=1,  label="Diffusion Volume")
+                    with gr.Row():
+                        n_iters = gr.Slider(100, 2000, value=500, step=100,
+                                            label="Training Iterations")
+                        lr_inp  = gr.Slider(1e-4, 1e-2, value=3e-4, step=1e-4,
+                                            label="Learning Rate")
+                    btn2    = gr.Button("🚀  Compress & Compare", variant="primary")
+                    stats2  = gr.Textbox(label="Results", lines=3, interactive=False)
+
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        orig_img  = gr.Image(label="📷 Original Slice",
+                                             type="numpy", height=380)
+                        recon_img = gr.Image(label="🤖 SIREN Reconstruction",
+                                             type="numpy", height=380)
+
+            btn2.click(compress_and_compare,
+                       inputs=[nifti_upload, sl2, vl2, n_iters, lr_inp],
+                       outputs=[orig_img, recon_img, stats2])
+
+        # ══════════════════════════════════════════════════════════════════════
+        # TAB 3 — About
+        # ══════════════════════════════════════════════════════════════════════
+        with gr.Tab("ℹ️ About"):
+            gr.Markdown(f"""
+## About this App
+
+**Physics-Informed SIREN MRI Compression** uses sinusoidal representation networks
+(SIRENs) to learn a compact neural implicit representation of diffusion MRI data.
+
+### How it works
+1. Each axial slice is represented by a small MLP with **sine activations** (SIREN)
+2. The network maps 2D spatial coordinates **(x, y) → signal intensities** across all diffusion volumes
+3. A **physics-informed loss** (Stejskal-Tanner constraint) regularises the network
+4. At inference time, coordinates are queried to reconstruct the full slice
+
+### Key advantages
+- 🗜️ **High compression ratio** — one small network per slice replaces raw voxel data
+- ⚡ **Resolution-agnostic** — can reconstruct at any spatial resolution
+- 🔬 **Physics-aware** — diffusion signal constraints improve anatomical fidelity
+- 🧩 **No codec artefacts** — continuous representation, no JPEG/JPEG2000 blocking
+
+### Model trained on
+[MGH-1010 Connectome Diffusion Microstructure Dataset](https://www.kaggle.com/datasets)
+
+| Property | Value |
+|---|---|
+| Architecture | `{config['model'].upper()}` |
+| Layers | `{config['num_layers']}` |
+| Hidden units | `{config['layer_size']}` |
+| w₀ | `{config['w0']}` |
+| Spatial slices | `{sz}` |
+| Diffusion volumes | `{vols}` |
+| Training data shape | `{sx} × {sy} × {sz} × {vols}` |
+
+### References
+- Sitzmann et al. (2020) — *Implicit Neural Representations with Periodic Activation Functions*
+- Stejskal & Tanner (1965) — *Spin diffusion measurements: spin echoes in the presence of a time-dependent field gradient*
+""")
 
 demo.launch()