Code updates

model.py → inference_brain2vec.py

@@ -1,9 +1,34 @@
-# model.py
-import os
-from typing import Optional
+#!/usr/bin/env python3
+
+"""
+inference_brain2vec.py
+
+Loads a pretrained Brain2vec VAE (AutoencoderKL) model and performs inference
+on one or more MRI images, generating reconstructions and latent parameters 
+(z_mu, z_sigma).
 
+Example usage:
+
+    # 1) Multiple file paths
+    python inference_brain2vec.py \
+        --checkpoint_path /path/to/autoencoder_checkpoint.pth \
+        --input_images /path/to/img1.nii.gz /path/to/img2.nii.gz \
+        --output_dir ./vae_inference_outputs \
+        --device cuda
+
+    # 2) Use a CSV containing image paths
+    python inference_brain2vec.py \
+        --checkpoint_path /path/to/autoencoder_checkpoint.pth \
+        --csv_input /path/to/images.csv \
+        --output_dir ./vae_inference_outputs
+"""
+
+import os
+import argparse
+import numpy as np
 import torch
 import torch.nn as nn
+from typing import Optional
 from monai.transforms import (
     Compose,
     CopyItemsD,
@@ -14,12 +39,12 @@ from monai.transforms import (
     ScaleIntensityD,
 )
 from generative.networks.nets import AutoencoderKL
+import pandas as pd
+
 
-# 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']),
@@ -29,15 +54,23 @@ transforms_fn = Compose([
     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.
+    a 5D tensor (batch=1, channel=1, D, H, W) for inference.
+
+    Args:
+        image_path (str): Path to the MRI (e.g. .nii.gz).
+        device (str): Device to place the tensor on.
+
+    Returns:
+        torch.Tensor: Shape (1, 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)
+    image_tensor = image_tensor.unsqueeze(0)  # => (1, 1, D, H, W)
     return image_tensor.to(device)
 
 
@@ -63,11 +96,11 @@ class Brain2vec(AutoencoderKL):
         Otherwise, return an uninitialized model.
 
         Args:
-            checkpoint_path (Optional[str]): path to a .pth checkpoint
+            checkpoint_path (Optional[str]): Path to a .pth checkpoint file.
             device (str): "cpu", "cuda", "mps", etc.
 
         Returns:
-            nn.Module: the loaded Brain2vec model on the chosen device
+            nn.Module: The loaded Brain2vec model on the chosen device.
         """
         model = Brain2vec(
             spatial_dims=3,
@@ -90,5 +123,101 @@ class Brain2vec(AutoencoderKL):
             model.load_state_dict(state_dict)
 
         model.to(device)
-        model.eval()  # ready for inference
-        return model
+        model.eval()
+        return model
+
+
+def main() -> None:
+    """
+    Main function to parse command-line arguments and run inference
+    with a pretrained Brain2vec model.
+    """
+    parser = argparse.ArgumentParser(
+        description="Inference script for a Brain2vec (VAE) model."
+    )
+    parser.add_argument(
+        "--checkpoint_path", type=str, required=True,
+        help="Path to the .pth checkpoint of the pretrained Brain2vec model."
+    )
+    parser.add_argument(
+        "--output_dir", type=str, default="./vae_inference_outputs",
+        help="Directory to save reconstructions and latent parameters."
+    )
+    parser.add_argument(
+        "--device", type=str, default="cpu",
+        help="Device to run inference on ('cpu', 'cuda', etc.)."
+    )
+    # Two ways to supply images: multiple file paths or a CSV
+    parser.add_argument(
+        "--input_images", type=str, nargs="*",
+        help="One or more MRI file paths (e.g. .nii.gz)."
+    )
+    parser.add_argument(
+        "--csv_input", type=str,
+        help="Path to a CSV file with an 'image_path' column."
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Load the model
+    model = Brain2vec.from_pretrained(
+        checkpoint_path=args.checkpoint_path,
+        device=args.device
+    )
+
+    # Gather image paths
+    if args.csv_input:
+        df = pd.read_csv(args.csv_input)
+        if "image_path" not in df.columns:
+            raise ValueError("CSV must contain a column named 'image_path'.")
+        image_paths = df["image_path"].tolist()
+    else:
+        if not args.input_images:
+            raise ValueError("Must provide either --csv_input or --input_images.")
+        image_paths = args.input_images
+
+    # Lists for stacking latent parameters later
+    all_z_mu = []
+    all_z_sigma = []
+
+    # Inference on each image
+    for i, img_path in enumerate(image_paths):
+        if not os.path.exists(img_path):
+            raise FileNotFoundError(f"Image not found: {img_path}")
+
+        print(f"[INFO] Processing image {i}: {img_path}")
+        img_tensor = preprocess_mri(img_path, device=args.device)
+
+        with torch.no_grad():
+            recon, z_mu, z_sigma = model.forward(img_tensor)
+
+        # Convert to NumPy
+        recon_np = recon.detach().cpu().numpy()  # shape: (1, 1, D, H, W)
+        z_mu_np = z_mu.detach().cpu().numpy()    # shape: (1, latent_channels, ...)
+        z_sigma_np = z_sigma.detach().cpu().numpy()
+
+        # Save each reconstruction (per image) as .npy
+        recon_path = os.path.join(args.output_dir, f"reconstruction_{i}.npy")
+        np.save(recon_path, recon_np)
+        print(f"[INFO] Saved reconstruction to {recon_path}")
+
+        # Store latent parameters for optional combined saving
+        all_z_mu.append(z_mu_np)
+        all_z_sigma.append(z_sigma_np)
+
+    # Combine latent parameters from all images and save
+    stacked_mu = np.concatenate(all_z_mu, axis=0)       # e.g., shape (N, latent_channels, ...)
+    stacked_sigma = np.concatenate(all_z_sigma, axis=0) # e.g., shape (N, latent_channels, ...)
+
+    mu_path = os.path.join(args.output_dir, "all_z_mu.npy")
+    sigma_path = os.path.join(args.output_dir, "all_z_sigma.npy")
+    np.save(mu_path, stacked_mu)
+    np.save(sigma_path, stacked_sigma)
+
+    print(f"[INFO] Saved z_mu of shape {stacked_mu.shape} to {mu_path}")
+    print(f"[INFO] Saved z_sigma of shape {stacked_sigma.shape} to {sigma_path}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -1,12 +1,15 @@
 # requirements.txt
 
-# PyTorch (CUDA or CPU version). For GPU install, see PyTorch docs for the correct wheel.
+# PyTorch (CUDA or CPU version).
 torch>=1.12
 
-# MONAI v1.2+ has the 'generative' subpackage with AutoencoderKL, PatchDiscriminator, etc.
-monai-weekly
+# Install MONAI Generative first
 monai-generative
 
+# Now force reinstall MONAI Weekly so its (newer) MONAI version takes precedence
+--force-reinstall
+monai-weekly
+
 # For perceptual losses in MONAI's generative module.
 lpips
 
@@ -17,4 +20,5 @@ nibabel
 tqdm
 tensorboard
 matplotlib
-datasets
+datasets
+scikit-learn

brain2vec.py → train_brain2vec.py

@@ -1,35 +1,20 @@
-# MIT License
-
-# Copyright (c) 2025
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# Forked from: https://github.com/LemuelPuglisi/BrLP
-
-# @inproceedings{puglisi2024enhancing,
-#   title={Enhancing spatiotemporal disease progression models via latent diffusion and prior knowledge},
-#   author={Puglisi, Lemuel and Alexander, Daniel C and Rav{\`\i}, Daniele},
-#   booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
-#   pages={173--183},
-#   year={2024},
-#   organization={Springer}
-# }
+#!/usr/bin/env python3
+
+"""
+train_brain2vec.py
+
+Trains a 3D VAE-based Brain2Vec model using MONAI. This script implements 
+autoencoder training with adversarial loss (via a patch discriminator),
+a perceptual loss, and KL divergence regularization for robust latent 
+representations. 
+
+Example usage:
+    python train_brain2vec.py train \
+        --dataset_csv /path/to/dataset.csv \
+        --cache_dir /path/to/cache \
+        --output_dir /path/to/output_dir \
+        --n_epochs 10
+"""
 
 import os
 os.environ["PYTORCH_WEIGHTS_ONLY"] = "False"
@@ -37,7 +22,6 @@ from typing import Optional, Union
 import pandas as pd
 import argparse
 import numpy as np
-
 import warnings
 import torch
 import torch.nn as nn
@@ -47,7 +31,6 @@ from torch.nn import L1Loss
 from torch.utils.data import DataLoader
 from torch.amp import autocast
 from torch.amp import GradScaler
-
 from generative.networks.nets import (
     AutoencoderKL, 
     PatchDiscriminator,
@@ -65,13 +48,11 @@ torch.serialization.add_safe_globals([_reconstruct])
 torch.serialization.add_safe_globals([MetaTensor])
 torch.serialization.add_safe_globals([ndarray])
 torch.serialization.add_safe_globals([dtype])
-
 from tqdm import tqdm
 import matplotlib.pyplot as plt
-
 from torch.utils.tensorboard import SummaryWriter
 
-# choosen resolution
+# voxel resolution
 RESOLUTION = 2                              
 
 # shape of the MNI152 (1mm^3) template
@@ -101,10 +82,7 @@ def load_if(checkpoints_path: Optional[str], network: nn.Module) -> nn.Module:
     """
     if checkpoints_path is not None:
         assert os.path.exists(checkpoints_path), 'Invalid path'
-        # Using context manager to allow MetaTensor
-        #with torch.serialization.safe_globals([MetaTensor]):
         network.load_state_dict(torch.load(checkpoints_path))
-        #network.load_state_dict(torch.load(checkpoints_path, map_location='cpu'))
     return network
 
 
@@ -140,7 +118,7 @@ def init_patch_discriminator(checkpoints_path: Optional[str] = None) -> nn.Modul
         checkpoints_path (Optional[str], optional): path of the checkpoints. Defaults to None.
 
     Returns:
-        nn.Module: the parch discriminator
+        nn.Module: the patch discriminator
     """
     patch_discriminator = PatchDiscriminator(spatial_dims=3, 
                                              num_layers_d=3, 
@@ -387,22 +365,6 @@ def train(
     train_df = dataset_df[dataset_df.split == 'train']
     trainset = get_dataset_from_pd(train_df, transforms_fn, cache_dir)
 
-    print(f"[DEBUG] Using cache_dir={cache_dir}")
-    print(f"[DEBUG] trainset length={len(trainset)}")
-
-    try:
-        sample_debug = trainset[0]  # Force a transform on the first record
-        print("[DEBUG] Successfully loaded sample 0 from trainset.")
-    except Exception as e:
-        print("[DEBUG] Error loading sample 0:", e)
-
-    import glob
-
-    hashfiles = glob.glob(os.path.join(cache_dir, "*.pt"))
-    print(f"[DEBUG] Found {len(hashfiles)} cached .pt files in {cache_dir}")
-    if hashfiles:
-        print("[DEBUG] Example cache file:", hashfiles[0])
-
     train_loader = DataLoader(
         dataset=trainset,
         num_workers=num_workers,
@@ -523,60 +485,11 @@ def train(
     print("Training completed and models saved.")
 
 
-def inference(
-    dataset_csv: str,
-    aekl_ckpt: str,
-    output_dir: str,
-    device: str = ('cuda' if torch.cuda.is_available() else 
-                   'cpu'),
-) -> None:
-    """
-    Perform inference to encode images into latent space.
-
-    Args:
-        dataset_csv (str): Path to the dataset CSV file.
-        aekl_ckpt (str): Path to the autoencoder checkpoint.
-        output_dir (str): Directory to save latent representations.
-        device (str, optional): Device to run the inference on. Defaults to 'cuda' if available.
-    """
-    DEVICE = device
-
-    autoencoder = init_autoencoder(aekl_ckpt).to(DEVICE).eval()
-
-    transforms_fn = transforms.Compose([
-        transforms.CopyItemsD(keys={'image_path'}, names=['image']),
-        transforms.LoadImageD(image_only=True, keys=['image']),
-        transforms.EnsureChannelFirstD(keys=['image']),
-        transforms.SpacingD(pixdim=RESOLUTION, keys=['image']),
-        transforms.ResizeWithPadOrCropD(spatial_size=INPUT_SHAPE_AE, mode='minimum', keys=['image']),
-        transforms.ScaleIntensityD(minv=0, maxv=1, keys=['image'])
-    ])
-
-    df = pd.read_csv(dataset_csv)
-
-    os.makedirs(output_dir, exist_ok=True)
-
-    with torch.no_grad():
-        for image_path in tqdm(df.image_path, total=len(df)):
-            destpath = os.path.join(
-                output_dir,
-                os.path.basename(image_path).replace('.nii.gz', '_embeddings.npz').replace('.nii', '_embeddings.npz')
-            )
-            if os.path.exists(destpath):
-                continue
-            mri_tensor = transforms_fn({'image_path': image_path})['image'].to(DEVICE)
-            mri_latent, _ = autoencoder.encode(mri_tensor.unsqueeze(0))
-            mri_latent = mri_latent.cpu().squeeze(0).numpy()
-            np.savez_compressed(destpath, data=mri_latent)
-
-    print("Inference completed and latent representations saved.")
-
-
 def main():
     """
-    Main function to parse command-line arguments and execute training or inference.
+    Main function to parse command-line arguments and execute training.
     """
-    parser = argparse.ArgumentParser(description="brain2vec Training and Inference Script")
+    parser = argparse.ArgumentParser(description="brain2vec Training Script")
 
     subparsers = parser.add_subparsers(dest='command', required=True, help='Sub-commands: train or infer')
 
@@ -594,12 +507,6 @@ def main():
     train_parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate.')
     train_parser.add_argument('--aug_p', type=float, default=0.8, help='Augmentation probability.')
 
-    # Inference Subparser
-    infer_parser = subparsers.add_parser('inference', help='Run inference to encode images.')
-    infer_parser.add_argument('--dataset_csv', type=str, required=True, help='Path to the dataset CSV file.')
-    infer_parser.add_argument('--aekl_ckpt', type=str, required=True, help='Path to the autoencoder checkpoint.')
-    infer_parser.add_argument('--output_dir', type=str, required=True, help='Directory to save latent representations.')
-
     args = parser.parse_args()
 
     if args.command == 'train':