README.md

---
license: apache-2.0
language:
  - en
task_categories:
  - image-classification
tags:
  - medical
  - brain-data
  - mri
pretty_name: 3D Brain Structure MRI Autoencoder
---

## 🧠 Model Summary
# brain2vec
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

# install Python libraries
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 from Hugging Face
python create_csv.py

mkdir ae_cache
mkdir ae_output

# train the model
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 ./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].

The model was trained with an effective batch size=16, 10 epochs, learning rate=1e-4 (see references 1 and 2).


# References
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,
  author    = {Jesse Brown and Clayton Young},
  title     = {Brain2vec: An Autoencoder Model for Brain Structure T1 MRIs},
  year      = {2025},
  url       = {https://huggingface.co/radiata-ai/brain2vec},
  note      = {Version 1.0},
  publisher = {Hugging Face}
}
```


# License
### 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.