anatomix (ICLR'25)

New 1: Added experimental 94M-parameter U-Net (anatomix-dev) and 26M-parameter 3D ViT (anatomix-dev-vit) models, both of which are trained on more and better data using better pretraining losses. Make sure to voxelwise normalize their features to have either unit norm or have zero mean with unit standard deviation prior to feature visualization or use in out-of-the-box feature-based applications like registration.

New 2: Interested in using anatomix features to enhance segmentation training regardless of the test domain and how much data you have? Check out MaskGen, a simple 5-line modification to standard training that uses anatomix to get SOTA domain generalization.

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.

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

Pull the model and weights from HuggingFace Hub:

from anatomix.model.load_from_hf import load_from_hf

model = load_from_hf("anatomix")          # 6M params. ICLR2025 version of the pretrained model

Or use one of the NEW experimental models and weights available here:

model = load_from_hf("anatomix-dev")        # 94M-parameter experimental UNet. More parameters, better features. Give it a shot.
# model = load_from_hf("anatomix-dev-vit")  # 26M-parameter experimental 3D ViT. Requires 128^3 inputs, but amenable to sliding window.

(NOTE: When visualizing features or when using these dev models for out-of-the-box applications like registration, make sure to voxelwise normalize the features across channels to have unit norm or zero mean with unit standard deviation, either will work.)

Or just load a local checkpoint:

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 = torch.compile(model)

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.

(If your task involves brains, abdomens, and generally interesting shapes, you might benefit from the two anatomix-dev models.)

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

conda create -n anatomix python=3.11
conda activate anatomix
git clone https://github.com/neel-dey/anatomix.git
cd anatomix
./install.sh   # auto-detects CUDA: cu126 default, cu130 for CUDA-13 drivers, cpu if no GPU

Or install the dependencies manually:

conda create -n anatomix python=3.11
conda activate anatomix
# change the cuda wheel to `cu130` if you're on CUDA13 or `cpu` if you're doing cpu-only
pip install numpy nibabel scipy scikit-image nilearn h5py matplotlib tensorboard tqdm monai torchio SimpleITK natsort huggingface_hub dynamic-network-architectures torch==2.13.0 torchvision==0.28.0 --extra-index-url https://download.pytorch.org/whl/cu126

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/                          # Pretrained model weights
│
├── pretraining/                            # Pretraining code and scripts 
│
├── synthetic-data-generation/              # Scripts to generate synthetic training data
│
├── anatomix/model/                         # Model definition and architecture
│
├── anatomix/registration/                  # Scripts for registration using the pretrained model
│
├── anatomix/segmentation/                  # Scripts for fine-tuning the model for semantic segmentation
│
└── README.md                               # This file

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 FireANTs
• Dataset-specific modeling details for paper

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}, 
}

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