In-Context Learning (ICL) for 3D Medical Imaging

March 14, 2026 · View on GitHub

ICCV 2025 Paper

This repository aims to advance the development of universal, In-Context Learning (ICL) based models for 3D medical imaging. It hosts the official implementation of Neuroverse3D, a versatile model for tasks like segmentation and image transformation without retraining. For a brief overview of the Neuroverse3D's capabilities, please see the introduction. It also include datasets generated by SynthICL, our method for generating synthetic datasets to enable training powerful ICL models from scratch.

Features

  • Memory-Efficient ICL Model: An ICL model specifically designed for 3D medical imaging that supports large context sizes without linear growth in memory consumption.
  • Pre-trained Checkpoint & Demo: Includes a pre-trained checkpoint for Neuroverse3D on neuroimaging data and a hands-on Jupyter Notebook (Demo.ipynb) to showcase its capabilities.
  • Training Pipeline: Provides the full source code for training the ICL segmentation model, along with a synthetic dataset generated via SynthICL, enabling researcher to train a capable 3D medical ICL model from scratch.

Getting Started

Follow the steps below to explore the model.

1. Environment Setup:

We offer three ways to set up the environment. We recommend using uv for the fastest dependency management.

  • Option 1: Using uv (Recommended) This project uses pyproject.toml for dependency management. You can set up the environment instantly using uv:

    # Install uv
pip install uv

# Sync dependencies (creates a virtual environment and installs packages)
uv sync

# Activate the environment
source .venv/bin/activate

  • Option 2: Using pip Ensure you have Python and PyTorch installed. You can install the dependencies listed in requirements.txt:

    pip install -r requirements.txt

  • Option 3: Using Docker Alternatively, you can directly download and use our provided Docker image, which contains all necessary environments.

2. Training (Optional)

If you wish to train the model from scratch, follow the steps below. To use our pre-trained model, you can skip to the next section.

First, download the provided synthetic dataset. We provide two download options:

You can easily add new datasets for training by following our dataset naming convention (consistent with nnUNet) and modifying train.py.

Then, run the two-stage training process:

# Stage 1: Initial training with a fixed context size
python train.py --train_gpus '(0)' \
                --precision 32 \
                --workers 8 \
                --model_name neuroverse3D \
                --nb_inner_channels 32 64 128 256 512 \
                --lr 0.00001 \
                --lr_decline_patience 100 \
                --data_dir /path/to/your/dataset/ \
                --skip_resize True \
                --context_size 3

# Stage 2: Fine-tuning with a random context size, loading the checkpoint from Stage 1
python train.py --train_gpus '(0)' \
                --checkpoint_path ./lightning_logs/version_0/checkpoints \
                --checkpoint_index -1 \
                --precision 32 \
                --workers 8 \
                --model_name neuroverse3D \
                --nb_inner_channels 32 64 128 256 512 \
                --lr 0.000002 \
                --lr_decline_patience 100 \
                --data_dir /path/to/your/dataset/ \
                --skip_resize True \
                --random_context_size True \
                --max_epochs 100

3. Inference

To run inference with the pre-trained model, first download the necessary files.

a) Download Checkpoint and Demo Images:

  • Pretrained Checkpoint: Download the Neuroverse3D checkpoint (neuroverse3D.ckpt) from Google Drive and place it in the ./checkpoint/ directory.
  • Demo Images: Download the demo images from Google Drive and place them in the ./Demo_data/ directory.
  • Jupyter Notebook: The Demo.ipynb notebook provides hands-on demonstrations of Neuroverse3D's capabilities.

b) Run the Model:

  • Run Model with Inference.py:
    This Inference.py script allows you to directly run Neuroverse3D on a given 3D medical imaging path.
    By default, this script assumes your data is organized in a format similar to nnUNet,
    i.e., structured as follows:

    Dataset001_BrainTumour/
 ├── imagesTr
 │   ├── BRATS_001_0000.nii.gz
 │   ├── BRATS_001_0001.nii.gz
 │   ├── BRATS_001_0002.nii.gz
 │   ├── BRATS_001_0003.nii.gz
 │   ├── BRATS_002_0000.nii.gz
 │   ├── BRATS_002_0001.nii.gz
 │   ├── BRATS_002_0002.nii.gz
 │   ├── BRATS_002_0003.nii.gz
 │   ├── ...
 └── labelsTr
     ├── BRATS_001.nii.gz
     ├── BRATS_002.nii.gz
     ├── ...

    Below are two usage examples with the provided Demo data:

    # For Segmentation Tasks
python Inference.py --checkpoint_path ./checkpoint/neuroverse3D.ckpt \
             --context_imgs Demo_data/seg/imgs \                 # Folder path for context images
             --context_imgs_modality 0000 \                     # Modality for context images
             --context_labs Demo_data/seg/labs \                # Folder path for context segmentation masks
             --target_imgs Demo_data/seg/imgs \                 # Folder path for target images
             --target_modality 0000 \                           # Modality for target images
             --target_output_path Demo_data/seg/preds \         # Folder to save predictions
             --task Seg                                         # Task type

# For Generation Tasks            
python Inference.py --checkpoint_path ./checkpoint/neuroverse3D.ckpt \
             --context_imgs Demo_data/mod_trans/imgs \          # Folder path for context images
             --context_imgs_modality 0000 \                     # Modality for context images
             --context_labs Demo_data/mod_trans/imgs \          # Folder path for context labels
             --context_labs_modality 0001 \                     # Modality for context labels
             --target_imgs Demo_data/mod_trans/imgs \           # Folder path for target images
             --target_modality 0000 \                           # Modality for target images
             --target_output_path Demo_data/mod_trans/preds \   # Folder to save predictions
             --task Gen                                         # Task type

  • Direct Model Execution:
    You can run the model directly using the following Python code:
    Ensure all input images are min-max normalized to [0, 1].

    from neuroverse3D.lightning_model import LightningModel
import torch
model = LightningModel.load_from_checkpoint(checkpoint_path)

# To perform a prediction (L = context size, spatial dimensions: H = W = D = 128)
with torch.no_grad():
    mask = model.forward(
        target_in,         # torch tensor (Batch, 1, H, W, D)
        context_in,        # torch tensor (Batch, L, 1, H, W, D)
        context_out,       # torch tensor (Batch, L, 1, H, W, D)
        gs=2,              # Mini-Context Size (positive integer). Smaller values reduce memory usage but decelerate processing.
        )  # -> (Batch, 1, H, W, D)

Citation

If you find Neuroverse3D model useful, please cite:

@inproceedings{hu2025neuroverse3d,
  title={Neuroverse3D: Developing In-Context Learning Universal Model for Neuroimaging in 3D},
  author={Hu, Jiesi and Peng, Hanyang and Yang, Yanwu and Guo, Xutao and Shang, Yang and Shi, Pengcheng and Ye, Chenfei and Ma, Ting},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={21721--21731},
  year={2025}
}

If find Synthetic Data useful, please cite:

@article{hu2025towards,
  title={Towards robust in-context learning for medical image segmentation via data synthesis},
  author={Hu, Jiesi and Yang, Yanwu and Ye, Zhiyu and Ye, Chenfei and Peng, Hanyang and Cao, Jianfeng and Ma, Ting},
  journal={arXiv preprint arXiv:2509.19711},
  year={2025}
}

Acknowledgements

This repository benefits from the excellent work provided by UniverSeg and Neuralizer. We extend our gratitude for their significant contributions to the field.