Code changes

README.md

@@ -13,23 +13,29 @@ pretty_name: 3D Brain Structure MRI Autoencoder
 
 ## 🧠 Model Summary
 # brain2vec
-An autoencoder model for brain structure T1 MRIs based on [Brain Latent Progression](https://github.com/LemuelPuglisi/BrLP/tree/main). The autoencoder takes in a 3d MRI NIfTI file and compresses to 1200 latent dimensions before reconstructing the image. The loss functions for training the autoencoder are:
+An autoencoder model for brain structure T1 MRIs (forked from [Brain Latent Progression](https://github.com/LemuelPuglisi/BrLP/tree/main)). The autoencoder takes in a 3d MRI NIfTI file and compresses to 1200 latent dimensions before reconstructing the image. The loss functions for training the autoencoder are:
 - [L1Loss](https://pytorch.org/docs/stable/generated/torch.nn.L1Loss.html)
 - [KLDivergenceLoss](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html)
 - [PatchAdversarialLoss](https://docs.monai.io/en/stable/losses.html#patchadversarialloss)
 - [PerceptualLoss](https://docs.monai.io/en/stable/losses.html#perceptualloss)
 
 
-
 # Training data
 [Radiata brain-structure](https://huggingface.co/datasets/radiata-ai/brain-structure): 3066 scans from 2085 individuals in the 'train' split. Mean age = 45.1 +- 24.5, including 2847 scans from cognitively normal subjects and 219 scans from individuals with an Alzheimer's disease clinical diagnosis.
 
+
 # Example usage
 ```
 # get brain2vec model repository
 git clone https://huggingface.co/radiata-ai/brain2vec
 cd brain2vec
 
+# pull pre-trained model weights
+sudo apt-get update
+sudo apt install git-lfs
+git lfs install
+git lfs pull
+
 # set up virtual environemt
 python3 -m venv venv_brain2vec
 source venv_brain2vec/bin/activate
@@ -38,54 +44,54 @@ source venv_brain2vec/bin/activate
 pip install -r requirements.txt
 
 # create the csv file inputs.csv listing the scan paths and other info
-# this script loads the radiata-ai/brain-structure dataset
+# this script loads the radiata-ai/brain-structure dataset from Hugging Face
 python create_csv.py
 
 mkdir ae_cache
 mkdir ae_output
 
-# install git lfs to pull large model weights
-sudo apt-get update
-sudo apt install git-lfs
-git lfs install
-git lfs pull
-
 # train the model
-nohup python brain2vec.py train \
-  --dataset_csv /home/ubuntu/brain2vec/inputs.csv \
+nohup python train_brain2vec.py \
+  --dataset_csv inputs.csv \
   --cache_dir   ./ae_cache \
   --output_dir  ./ae_output \
   --n_epochs    10 \
 > train_log.txt 2>&1 &
 
 # model inference
+# for a set of scans in inputs.csv
+python inference_brain2vec.py \
+  --checkpoint_path /path/to/model.pth \
+  --csv_input inputs.csv \
+  --output_dir ./ae_output \
+  --embeddings_filename ae_embeddings_all.npy
+
+# or for individual scans
 python inference_brain2vec.py \
   --checkpoint_path /path/to/model.pth \
   --input_images /path/to/img1.nii.gz /path/to/img2.nii.gz \
-  --output_dir ./vae_inference_outputs \
-  --embeddings_filename pca_output/pca_embeddings_2.npy \
-  --save_recons
+  --output_dir ./ae_output \
+  --embeddings_filename ae_embeddings_2.npy
 ```
 
 # Methods
 Input scan image dimensions are 113x137x113, 1.5mm^3 resolution, aligned to MNI152 space (see [radiata-ai/brain-structure](https://huggingface.co/datasets/radiata-ai/brain-structure)).  
 
-The image transform crops to 80 x 96 x 80, 2mm^3 resolution, and scales image intensity to range [0,1]. Images are flattened to 614400-length 1D vectors.  
+The image transform crops to 80 x 96 x 80, 2mm^3 resolution, and scales image intensity to range [0,1].
+
+The model was trained with an effective batch size=16, 10 epochs, learning rate=1e-4 (see references 1 and 2).
 
-10 epochs
-    max_batch_size: int = 2,
-    batch_size: int = 16,
-    lr: float = 1e-4,
 
 # References
-Puglisi
-Pinaya
+1. Puglisi L, Alexander DC, Ravì D. Enhancing Spatiotemporal Disease Progression Models via Latent Diffusion and Prior Knowledge [Internet]. arXiv; 2024. Available from: http://arxiv.org/abs/2405.03328
+2. Pinaya WHL, Tudosiu PD, Dafflon J, Costa PF da, Fernandez V, Nachev P, et al. Brain Imaging Generation with Latent Diffusion Models [Internet]. arXiv; 2022. Available from: http://arxiv.org/abs/2209.07162
+
 
 # Citation
 ```
-@misc{Radiata-Brain2Vec,
+@misc{Radiata-Brain2vec,
   author    = {Jesse Brown and Clayton Young},
-  title     = {brain2vec_PCA: A VAE Model for Brain Structure T1 MRIs},
+  title     = {Brain2vec: An Autoencoder Model for Brain Structure T1 MRIs},
   year      = {2025},
   url       = {https://huggingface.co/radiata-ai/brain2vec},
   note      = {Version 1.0},
@@ -93,25 +99,20 @@ Pinaya
 }
 ```
 
+
 # License
-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.
+### Apache License 2.0
+
+Copyright 2025 Jesse Brown
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at:
+
+[http://www.apache.org/licenses/LICENSE-2.0](http://www.apache.org/licenses/LICENSE-2.0)
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

inference_brain2vec.py

@@ -143,10 +143,6 @@ def main() -> None:
         "--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="*",
@@ -172,10 +168,13 @@ def main() -> None:
 
     os.makedirs(args.output_dir, exist_ok=True)
 
-    # Load the model
+    # After parsing args, add:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Then pass that device to the model:
     model = Brain2vec.from_pretrained(
         checkpoint_path=args.checkpoint_path,
-        device=args.device
+        device=device
     )
 
     # Gather image paths
@@ -199,7 +198,7 @@ def main() -> None:
             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)
+        img_tensor = preprocess_mri(img_path, device=device)
 
         with torch.no_grad():
             recon, z_mu, z_sigma = model.forward(img_tensor)

train_brain2vec.py

@@ -9,10 +9,10 @@ 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 \
+    python train_brain2vec.py \
+        --dataset_csv inputs.csv \
+        --cache_dir ./ae_cache \
+        --output_dir ./ae_output \
         --n_epochs 10
 """
 
@@ -487,50 +487,39 @@ def train(
 
 def main():
     """
-    Main function to parse command-line arguments and execute training.
+    Main function to parse command-line arguments and run train().
     """
+    import argparse
+
     parser = argparse.ArgumentParser(description="brain2vec Training Script")
 
-    subparsers = parser.add_subparsers(dest='command', required=True, help='Sub-commands: train or infer')
-
-    # Training Subparser
-    train_parser = subparsers.add_parser('train', help='Train the models.')
-    train_parser.add_argument('--dataset_csv', type=str, required=True, help='Path to the dataset CSV file.')
-    train_parser.add_argument('--cache_dir', type=str, required=True, help='Directory for caching data.')
-    train_parser.add_argument('--output_dir', type=str, required=True, help='Directory to save model checkpoints.')
-    train_parser.add_argument('--aekl_ckpt', type=str, default=None, help='Path to the autoencoder checkpoint.')
-    train_parser.add_argument('--disc_ckpt', type=str, default=None, help='Path to the discriminator checkpoint.')
-    train_parser.add_argument('--num_workers', type=int, default=8, help='Number of data loader workers.')
-    train_parser.add_argument('--n_epochs', type=int, default=5, help='Number of training epochs.')
-    train_parser.add_argument('--max_batch_size', type=int, default=2, help='Actual batch size per iteration.')
-    train_parser.add_argument('--batch_size', type=int, default=16, help='Expected (effective) batch size.')
-    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.')
+    parser.add_argument('--dataset_csv', type=str, required=True, help='Path to the dataset CSV file.')
+    parser.add_argument('--cache_dir', type=str, required=True, help='Directory for caching data.')
+    parser.add_argument('--output_dir', type=str, required=True, help='Directory to save model checkpoints.')
+    parser.add_argument('--aekl_ckpt', type=str, default=None, help='Path to the autoencoder checkpoint.')
+    parser.add_argument('--disc_ckpt', type=str, default=None, help='Path to the discriminator checkpoint.')
+    parser.add_argument('--num_workers', type=int, default=8, help='Number of data loader workers.')
+    parser.add_argument('--n_epochs', type=int, default=5, help='Number of training epochs.')
+    parser.add_argument('--max_batch_size', type=int, default=2, help='Actual batch size per iteration.')
+    parser.add_argument('--batch_size', type=int, default=16, help='Expected (effective) batch size.')
+    parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate.')
+    parser.add_argument('--aug_p', type=float, default=0.8, help='Augmentation probability.')
 
     args = parser.parse_args()
 
-    if args.command == 'train':
-        train(
-            dataset_csv=args.dataset_csv,
-            cache_dir=args.cache_dir,
-            output_dir=args.output_dir,
-            aekl_ckpt=args.aekl_ckpt,
-            disc_ckpt=args.disc_ckpt,
-            num_workers=args.num_workers,
-            n_epochs=args.n_epochs,
-            max_batch_size=args.max_batch_size,
-            batch_size=args.batch_size,
-            lr=args.lr,
-            aug_p=args.aug_p,
-        )
-    elif args.command == 'infer':
-        inference(
-            dataset_csv=args.dataset_csv,
-            aekl_ckpt=args.aekl_ckpt,
-            output_dir=args.output_dir,
-        )
-    else:
-        parser.print_help()
+    train(
+        dataset_csv=args.dataset_csv,
+        cache_dir=args.cache_dir,
+        output_dir=args.output_dir,
+        aekl_ckpt=args.aekl_ckpt,
+        disc_ckpt=args.disc_ckpt,
+        num_workers=args.num_workers,
+        n_epochs=args.n_epochs,
+        max_batch_size=args.max_batch_size,
+        batch_size=args.batch_size,
+        lr=args.lr,
+        aug_p=args.aug_p,
+    )
 
 
 if __name__ == '__main__':