Skip to content

[ICLR 2025] Learning General-purpose Biomedical Volume Representations using Randomized Synthesis

License

Notifications You must be signed in to change notification settings

neel-dey/anatomix

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

23 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Paper shield License: MIT Colab tutorial 1

anatomix (ICLR'25)

Highlight collage

anatomix is a general-purpose feature extractor for 3D volumes. For any new biomedical dataset or task,

  • Its out-of-the-box features are invariant to most forms of nuisance imaging variation.
  • Its out-of-the-box weights are a good initialization for finetuning when given limited annotations.

Without any dataset or domain specific pretraining, this respectively leads to:

  • SOTA 3D training-free multi-modality image registration
  • SOTA 3D few-shot semantic segmentation

How? It's weights were contrastively pretrained on wildly variable synthetic volumes to learn approximate appearance invariance and pose equivariance for images with randomly sampled biomedical shape configurations with random intensities and artifacts.

Load weights for inference / feature extraction / finetuning

anatomix is just a pretrained UNet! Use it for whatever you like.

import torch
from anatomix.model.network import Unet

model = Unet(
    dimension=3,  # Only 3D supported for now
    input_nc=1,  # number of input channels
    output_nc=16,  # number of output channels
    num_downs=4,  # number of downsampling layers
    ngf=16,  # channel multiplier
)
model.load_state_dict(
    torch.load("./model-weights/anatomix.pth"),
    strict=True,
)

See how to use it on real data for feature extraction or registration in this tutorial. Or if you want to finetune for your own task, check out this tutorial instead.

Install dependencies:

All scripts will require the anatomix environment defined below to run.

conda create -n anatomix python=3.9
conda activate anatomix
git clone https://github.com/neel-dey/anatomix.git
cd anatomix
pip install -e .

Or install the dependencies manually:

conda create -n anatomix python=3.9
conda activate anatomix
pip install numpy==1.24.1 nibabel scipy scikit-image nilearn h5py matplotlib torch==2.0.1 tensorboard tqdm monai==1.3.2

Folder organization

This repository contains:

  • Model weights for anatomix (and its variants)
  • Scripts for 3D nonrigid registration using anatomix features
  • Scripts for finetuning anatomix weights for semantic segmentation.

Each subfolder (described below) has its own README to detail its use.

root-directory/
│
├── model-weights/                          # Pre-trained model weights
│
├── anatomix/model/                         # Model definition and architecture
│
├── anatomix/synthetic-data-generation/     # Scripts to generate synthetic training data
│
├── anatomix/registration/                  # Scripts for registration using the pretrained model
│
├── anatomix/segmentation/                  # Scripts for fine-tuning the model for semantic segmentation
│
└── README.md                               # This file

Roadmap

The current repository is just an initial push. It will be refactored and some quality-of-life and core aspects will be pushed as well in the coming weeks. These include:

  • Contrastive pretraining code
  • Colab 3D feature extraction tutorial
  • Colab 3D multimodality registration tutorial
  • Colab 3D few-shot finetuning tutorial
  • General-purpose registration interface with ANTs
  • Dataset-specific modeling details for paper

Citation

If you find our work useful, please consider citing:

@misc{dey2024learninggeneralpurposebiomedicalvolume,
      title={Learning General-Purpose Biomedical Volume Representations using Randomized Synthesis}, 
      author={Neel Dey and Benjamin Billot and Hallee E. Wong and Clinton J. Wang and Mengwei Ren and P. Ellen Grant and Adrian V. Dalca and Polina Golland},
      year={2024},
      eprint={2411.02372},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2411.02372}, 
}

Acknowledgements

Portions of this repository have been taken from the Contrastive Unpaired Translation and ConvexAdam repositories and modified. Thanks!

License

We use the standard MIT license.

Releases

No releases published

Packages

No packages published