Update app.py

app.py

@@ -1,164 +1,183 @@
 import os
-import gradio as gr
-import torch
 import pickle
 import numpy as np
-import nibabel as nib
-
-from siren import Siren
-
-
-# =========================
-# DEVICE
-# =========================
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-# =========================
-# LOAD MODEL
-# =========================
-checkpoint = torch.load(
-    "sirenMRI_full_model_final.pt",
-    map_location=device
+import torch
+import torch.nn as nn
+from math import sqrt
+import streamlit as st
+
+# ── Streamlit page config ────────────────────────────────────────────────────
+st.set_page_config(
+    page_title="Physics-Informed SIREN MRI Compression",
+    page_icon="🧠",
+    layout="wide",
 )
 
-config = checkpoint["config"]
-
-model = 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"]
-).to(device)
-
-model.load_state_dict(checkpoint["model_state_dict"])
-model.eval()
-
-
-# =========================
-# LOAD SCALERS
-# =========================
-with open("scalers.pkl", "rb") as f:
-    scalers = pickle.load(f)
-
-
-# =========================
-# CREATE COORDINATES
-# =========================
-def create_coordinates(h, w):
-
-    y_coords = np.linspace(-1, 1, h)
-    x_coords = np.linspace(-1, 1, w)
-
-    yy, xx = np.meshgrid(y_coords, x_coords, indexing='ij')
-
-    coords = np.stack([yy, xx], axis=-1)
-
-    coords = coords.reshape(-1, 2)
-
-    return torch.tensor(coords, dtype=torch.float32)
-
-
-# =========================
-# RECONSTRUCTION FUNCTION
-# =========================
-def reconstruct(nifti_file):
-
-    try:
-
-        # =========================
-        # LOAD MRI
-        # =========================
-        nii = nib.load(nifti_file.name)
-        img = nii.get_fdata()
-
-        original_shape = img.shape
-
-        # =========================
-        # HANDLE 4D MRI
-        # =========================
-        if len(img.shape) != 4:
-            return "❌ Expected a 4D MRI (.nii or .nii.gz)", None
-
-        h, w, slices, vols = img.shape
+# ══════════════════════════════════════════════════════════════════════════════
+# 1.  Model definitions — must match exactly what was used during training
+# ══════════════════════════════════════════════════════════════════════════════
+
+class Sine(nn.Module):
+    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):
+        super().__init__()
+        self.dim_in = dim_in
+        self.is_first = is_first
+        self.linear = nn.Linear(dim_in, dim_out, bias=use_bias)
+        w_std = (1 / dim_in) if is_first else (sqrt(c / dim_in) / w0)
+        nn.init.uniform_(self.linear.weight, -w_std, w_std)
+        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):
+        super().__init__()
+        layers = []
+        for i in range(num_layers):
+            is_first = i == 0
+            layer_w0 = w0_initial if is_first else w0
+            layer_dim_in = dim_in if is_first else dim_hidden
+            layers.append(SirenLayer(
+                dim_in=layer_dim_in, dim_out=dim_hidden,
+                w0=layer_w0, use_bias=use_bias, is_first=is_first,
+            ))
+        self.net = nn.Sequential(*layers)
+        final_activation = nn.Identity() if final_activation is None else final_activation
+        self.last_layer = SirenLayer(
+            dim_in=dim_hidden, dim_out=dim_out,
+            w0=w0, use_bias=use_bias, activation=final_activation,
+        )
 
-        reconstructed = np.zeros_like(img)
+    def forward(self, x):
+        return self.last_layer(self.net(x))
 
-        # =========================
-        # GENERATE COORDINATES
-        # =========================
-        coords = create_coordinates(h, w).to(device)
 
-        # =========================
-        # RECONSTRUCT EACH SLICE
-        # =========================
-        with torch.no_grad():
+class SirenMRIModel(nn.Module):
+    """Multi-slice wrapper — one Siren per axial slice."""
 
-            for s in range(slices):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        self.models = nn.ModuleList()
+        for _ in range(config["sz"]):
+            self.models.append(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"],
+            ))
 
-                pred = model(coords)
+    def forward(self, coords: torch.Tensor, slice_idx: int) -> torch.Tensor:
+        return self.models[slice_idx](coords)
 
-                pred = pred.cpu().numpy()
 
-                # inverse normalization
-                pred = scalers[s].inverse_transform(pred)
+# ══════════════════════════════════════════════════════════════════════════════
+# 2.  Load model + scalers  (cached so they only load once)
+# ══════════════════════════════════════════════════════════════════════════════
 
-                pred = pred.reshape(h, w, vols)
+@st.cache_resource
+def load_model():
+    model_path   = "sirenMRI_full_model_final.pt"
+    scalers_path = "scalers.pkl"
 
-                reconstructed[:, :, s, :] = pred
+    if not os.path.exists(model_path):
+        st.error(f"Model file not found: {model_path}")
+        st.stop()
+    if not os.path.exists(scalers_path):
+        st.error(f"Scalers file not found: {scalers_path}")
+        st.stop()
 
-        # =========================
-        # SAVE OUTPUT MRI
-        # =========================
-        output_path = "reconstructed_output.nii.gz"
+    checkpoint = torch.load(model_path, map_location="cpu")
+    config     = checkpoint["config"]
 
-        recon_nii = nib.Nifti1Image(
-            reconstructed,
-            affine=nii.affine,
-            header=nii.header
-        )
+    # ── THE FIX: build SirenMRIModel, not bare Siren ──────────────────────
+    model = SirenMRIModel(config)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
 
-        nib.save(recon_nii, output_path)
+    with open(scalers_path, "rb") as f:
+        scalers = pickle.load(f)
 
-        return (
-            f"""✅ Reconstruction complete
+    return model, scalers, config, checkpoint["input_shape"]
 
-Original Shape: {original_shape}
 
-Reconstructed MRI saved successfully.
-""",
-            output_path
-        )
+model, scalers, config, input_shape = load_model()
+sx, sy, sz, vols = input_shape
 
-    except Exception as e:
-        return f"❌ Error: {str(e)}", None
 
+# ══════════════════════════════════════════════════════════════════════════════
+# 3.  Inference helper
+# ══════════════════════════════════════════════════════════════════════════════
 
-# =========================
-# GRADIO UI
-# =========================
-interface = gr.Interface(
-    fn=reconstruct,
+def reconstruct_slice(slice_idx: int) -> np.ndarray:
+    """Return reconstructed slice as (sx, sy, vols) float32 array."""
+    # Build normalised (x, y) coordinate grid in [-1, 1]
+    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)  # (sx*sy, 2)
 
-    inputs=gr.File(
-        label="Upload MRI (.nii or .nii.gz)"
-    ),
+    with torch.no_grad():
+        pred = model(coords, slice_idx).numpy()  # (sx*sy, vols)
 
-    outputs=[
-        gr.Textbox(label="Status"),
-        gr.File(label="Download Reconstructed MRI")
-    ],
+    # Inverse-transform back to original intensity range
+    pred = scalers[slice_idx].inverse_transform(pred)
+    return pred.reshape(sx, sy, vols)
 
-    title="Physics-Informed SIREN MRI Reconstruction",
 
-    description="""
-Upload a 4D MRI scan (.nii or .nii.gz).
+# ══════════════════════════════════════════════════════════════════════════════
+# 4.  Streamlit UI
+# ══════════════════════════════════════════════════════════════════════════════
 
-The model reconstructs the MRI using a trained
-Physics-Informed SIREN neural representation.
-"""
+st.title("🧠 Physics-Informed SIREN MRI Compression")
+st.markdown(
+    "Reconstruct axial slices from the compressed SIREN representation. "
+    "Use the sliders to choose a slice and a diffusion volume."
 )
 
-interface.launch()
+col1, col2 = st.columns(2)
+with col1:
+    slice_idx = st.slider("Axial slice", 0, sz - 1, sz // 2)
+with col2:
+    vol_idx = st.slider("Diffusion volume (b-value index)", 0, vols - 1, 0)
+
+if st.button("Reconstruct slice"):
+    with st.spinner("Running SIREN inference..."):
+        recon = reconstruct_slice(slice_idx)   # (sx, sy, vols)
+        img   = recon[:, :, vol_idx]
+
+    # Normalise to [0, 1] for display
+    img_norm = (img - img.min()) / (img.max() - img.min() + 1e-8)
+
+    st.image(
+        img_norm,
+        caption=f"Reconstructed slice {slice_idx}, volume {vol_idx}",
+        use_column_width=True,
+        clamp=True,
+    )
+
+    st.success(f"Slice shape: {img.shape} | Value range: [{img.min():.3f}, {img.max():.3f}]")
+
+st.divider()
+st.caption(
+    f"Model config — layers: {config['num_layers']}, "
+    f"hidden size: {config['layer_size']}, "
+    f"w0: {config['w0']}, slices: {sz}, volumes: {vols}"
+)