Updates

brain2vec.py

@@ -72,7 +72,7 @@ import matplotlib.pyplot as plt
 from torch.utils.tensorboard import SummaryWriter
 
 # choosen resolution
-RESOLUTION = 1.5                                
+RESOLUTION = 2                              
 
 # shape of the MNI152 (1mm^3) template
 INPUT_SHAPE_1mm = (182, 218, 182)   
@@ -616,7 +616,7 @@ def main():
             lr=args.lr,
             aug_p=args.aug_p,
         )
-    elif args.command == 'inference':
+    elif args.command == 'infer':
         inference(
             dataset_csv=args.dataset_csv,
             aekl_ckpt=args.aekl_ckpt,

model.py

@@ -0,0 +1,94 @@
+# model.py
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from monai.transforms import (
+    Compose,
+    CopyItemsD,
+    LoadImageD,
+    EnsureChannelFirstD,
+    SpacingD,
+    ResizeWithPadOrCropD,
+    ScaleIntensityD,
+)
+from generative.networks.nets import AutoencoderKL
+
+# Constants for your typical config
+RESOLUTION = 2
+INPUT_SHAPE_AE = (80, 96, 80)
+
+# Define the exact transform pipeline for input MRI
+transforms_fn = Compose([
+    CopyItemsD(keys={'image_path'}, names=['image']),
+    LoadImageD(image_only=True, keys=['image']),
+    EnsureChannelFirstD(keys=['image']),
+    SpacingD(pixdim=RESOLUTION, keys=['image']),
+    ResizeWithPadOrCropD(spatial_size=INPUT_SHAPE_AE, mode='minimum', keys=['image']),
+    ScaleIntensityD(minv=0, maxv=1, keys=['image']),
+])
+
+def preprocess_mri(image_path: str, device: str = "cpu") -> torch.Tensor:
+    """
+    Preprocess an MRI using MONAI transforms to produce
+    a 5D tensor (batch=1, channels=1, D, H, W) for inference.
+    """
+    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)
+
+
+class Brain2vec(AutoencoderKL):
+    """
+    Subclass of MONAI's AutoencoderKL that includes:
+      - a from_pretrained(...) for loading a .pth checkpoint
+      - uses the existing forward(...) that returns (reconstruction, z_mu, z_sigma)
+
+    Usage:
+      >>> model = Brain2vec.from_pretrained("my_checkpoint.pth", device="cuda")
+      >>> image_tensor = preprocess_mri("/path/to/mri.nii.gz", device="cuda")
+      >>> reconstruction, z_mu, z_sigma = model.forward(image_tensor)
+    """
+
+    @staticmethod
+    def from_pretrained(
+        checkpoint_path: Optional[str] = None,
+        device: str = "cpu"
+    ) -> nn.Module:
+        """
+        Load a pretrained Brain2vec (AutoencoderKL) if a checkpoint_path is provided.
+        Otherwise, return an uninitialized model.
+
+        Args:
+            checkpoint_path (Optional[str]): path to a .pth checkpoint
+            device (str): "cpu", "cuda", "mps", etc.
+
+        Returns:
+            nn.Module: the loaded Brain2vec model on the chosen device
+        """
+        model = Brain2vec(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=1,
+            latent_channels=1,
+            num_channels=(64, 128, 128, 128),
+            num_res_blocks=2,
+            norm_num_groups=32,
+            norm_eps=1e-06,
+            attention_levels=(False, False, False, False),
+            with_decoder_nonlocal_attn=False,
+            with_encoder_nonlocal_attn=False,
+        )
+
+        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.to(device)
+        model.eval()  # ready for inference
+        return model