Updated model.py and added .joblib

model.py

@@ -1,9 +1,9 @@
 # model.py
 import os
-from typing import Optional
-
+import numpy as np
 import torch
 import torch.nn as nn
+
 from monai.transforms import (
     Compose,
     CopyItemsD,
@@ -14,11 +14,20 @@ from monai.transforms import (
     ScaleIntensityD,
 )
 
-# Constants for your typical config
+# If you used joblib or pickle to save your PCA model:
+from joblib import load  # or "import pickle"
+
+#################################################
+# Constants
+#################################################
 RESOLUTION = 2
-INPUT_SHAPE_AE = (80, 96, 80)
+INPUT_SHAPE_AE = (80, 96, 80)  # The typical shape from your pipelines
+FLATTENED_DIM = INPUT_SHAPE_AE[0] * INPUT_SHAPE_AE[1] * INPUT_SHAPE_AE[2]
 
-# Define the exact transform pipeline for input MRI
+
+#################################################
+# Define MONAI Transforms for Preprocessing
+#################################################
 transforms_fn = Compose([
     CopyItemsD(keys={'image_path'}, names=['image']),
     LoadImageD(image_only=True, keys=['image']),
@@ -28,94 +37,79 @@ transforms_fn = Compose([
     ScaleIntensityD(minv=0, maxv=1, keys=['image']),
 ])
 
-def preprocess_mri(image_path: str, device: str = "cpu") -> torch.Tensor:
+
+def preprocess_mri(image_path: str) -> torch.Tensor:
     """
     Preprocess an MRI using MONAI transforms to produce
-    a 5D tensor (batch=1, channels=1, D, H, W) for inference.
+    a 5D Torch tensor: (batch=1, channel=1, D, H, W).
     """
     data_dict = {"image_path": image_path}
     output_dict = transforms_fn(data_dict)
-    image_tensor = output_dict["image"]  # shape: (1, D, H, W)
-    image_tensor = image_tensor.unsqueeze(0)  # => (batch=1, channel=1, D, H, W)
-    return image_tensor.to(device)
+    # shape => (1, D, H, W)
+    image_tensor = output_dict["image"].unsqueeze(0)  # => (batch=1, channel=1, D, H, W)
+    return image_tensor.float()  # typically float32
 
 
-class ShallowLinearAutoencoder(nn.Module):
+#################################################
+# PCA "Autoencoder" Wrapper
+#################################################
+class PCABrain2vec(nn.Module):
     """
-    A purely linear autoencoder with one hidden layer.
-    - Flatten input into a vector
-    - Linear encoder (no activation)
-    - Linear decoder (no activation)
-    - Reshape output to original volume shape
-    """
-    def __init__(self, input_shape=(80, 96, 80), hidden_size=1200):
-        super().__init__()
-        self.input_shape = input_shape
-        self.input_dim = input_shape[0] * input_shape[1] * input_shape[2]
-        self.hidden_size = hidden_size
-
-        # Encoder (no activation for PCA-like behavior)
-        self.encoder = nn.Sequential(
-            nn.Flatten(),
-            nn.Linear(self.input_dim, self.hidden_size),
-        )
-
-        # Decoder (no activation)
-        self.decoder = nn.Sequential(
-            nn.Linear(self.hidden_size, self.input_dim),
-        )
+    A PCA-based 'autoencoder' that mimics the old interface:
+      - from_pretrained(...) to load a PCA model from disk
+      - forward(...) returns (reconstruction, embedding, None)
 
-    def encode(self, x: torch.Tensor):
-        return self.encoder(x)
+    Under the hood, it:
+      - takes in a torch tensor shape (N, 1, D, H, W)
+      - flattens it (N, 614400)
+      - uses PCA's transform(...) to get embeddings => shape (N, n_components)
+      - uses inverse_transform(...) to get reconstructions => shape (N, 614400)
+      - reshapes back to (N, 1, D, H, W)
+    """
 
-    def decode(self, z: torch.Tensor):
-        out = self.decoder(z)
-        # Reshape to (N, 1, D, H, W)
-        return out.view(-1, 1, *self.input_shape)
+    def __init__(self, pca_model=None):
+        super().__init__()
+        # We'll store the fitted PCA model (from scikit-learn)
+        self.pca_model = pca_model  # e.g., an instance of IncrementalPCA or PCA
 
     def forward(self, x: torch.Tensor):
         """
-        Return (reconstruction, embedding, None) to keep a similar API
-        to the old VAE-based code, though there's no σ for sampling.
+        Returns (reconstruction, embedding, None).
+
+        1) Convert x => numpy array => flatten => (N, 614400)
+        2) embedding = pca_model.transform(flat_x)
+        3) reconstruction_np = pca_model.inverse_transform(embedding)
+        4) reshape => (N, 1, 80, 96, 80)
+        5) convert to torch => return (recon, embed, None)
         """
-        z = self.encode(x)
-        reconstruction = self.decode(z)
-        return reconstruction, z, None
+        # Expect x shape => (N, 1, D, H, W) => flatten to (N, D*H*W)
+        n_samples = x.shape[0]
+        # Convert to CPU np
+        x_cpu = x.detach().cpu().numpy()  # shape: (N, 1, D, H, W)
+        x_flat = x_cpu.reshape(n_samples, -1)  # shape: (N, 614400)
 
+        # PCA transform => embeddings shape (N, n_components)
+        embedding_np = self.pca_model.transform(x_flat)
 
-class Brain2vec(nn.Module):
-    """
-    A wrapper around the ShallowLinearAutoencoder, providing a from_pretrained(...)
-    method for model loading, mirroring the old usage with AutoencoderKL.
-    """
-    def __init__(self, device: str = "cpu"):
-        super().__init__()
-        # Instantiate the shallow linear model
-        self.model = ShallowLinearAutoencoder(input_shape=INPUT_SHAPE_AE, hidden_size=1200)
-        self.to(device)
+        # PCA inverse_transform => recon shape (N, 614400)
+        recon_np = self.pca_model.inverse_transform(embedding_np)
+        # Reshape back => (N, 1, 80, 96, 80)
+        recon_np = recon_np.reshape(n_samples, 1, *INPUT_SHAPE_AE)
 
-    def forward(self, x: torch.Tensor):
-        """
-        Forward pass that returns (reconstruction, embedding, None).
-        """
-        return self.model(x)
+        # Convert back to torch
+        reconstruction_torch = torch.from_numpy(recon_np).float()
+        embedding_torch = torch.from_numpy(embedding_np).float()
+        return reconstruction_torch, embedding_torch, None
 
     @staticmethod
-    def from_pretrained(
-        checkpoint_path: Optional[str] = None,
-        device: str = "cpu"
-    ) -> nn.Module:
+    def from_pretrained(pca_path: str):
         """
-        Load a pretrained ShallowLinearAutoencoder if a checkpoint path is provided.
-        Args:
-            checkpoint_path (Optional[str]): path to a .pth checkpoint
-            device (str): "cpu", "cuda", etc.
+        Load a pre-trained PCA model (pickled or joblib).
+        Returns an instance of PCABrain2vec with that model.
         """
-        model = Brain2vec(device=device)
-        if checkpoint_path is not None:
-            if not os.path.exists(checkpoint_path):
-                raise FileNotFoundError(f"Checkpoint {checkpoint_path} not found.")
-            state_dict = torch.load(checkpoint_path, map_location=device)
-            model.load_state_dict(state_dict)
-        model.eval()
-        return model
+        if not os.path.exists(pca_path):
+            raise FileNotFoundError(f"Could not find PCA model at {pca_path}")
+        # Example: pca_model = pickle.load(open(pca_path, 'rb'))
+        # or use joblib:
+        pca_model = load(pca_path)
+        return PCABrain2vec(pca_model=pca_model)

pca_model.joblib

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