# In-Context Learning (ICL) for 3D Medical Imaging [![ICCV 2025](https://img.shields.io/badge/ICCV-2025-blue.svg)](https://arxiv.org/pdf/2503.02410v1) [![Paper](https://img.shields.io/badge/arxiv-2503.02410-b31b1b.svg)](https://arxiv.org/abs/2509.19711) 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](https://arxiv.org/pdf/2503.02410v1)**, 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](neuroverse3D/intro.md). It also include datasets generated by **[SynthICL](https://arxiv.org/abs/2509.19711)**, 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](https://github.com/astral-sh/uv): ```bash # 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`: ```bash pip install -r requirements.txt ``` * **Option 3: Using Docker** Alternatively, you can directly download and use our provided [Docker image](https://drive.google.com/file/d/1YuG8xlVCeLsUM0KGLbp8Xe_rBYYLfYtF/view?usp=sharing), 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: * **Google Drive:** [Download Link](https://drive.google.com/drive/folders/1TRLx_JRoHZdo81kJt9rAL3oQ6WoIvl-c?usp=sharing) * **Baidu Netdisk:** [Download Link](https://pan.baidu.com/s/1GNowsAfZE2vVIo1tW4O6vQ) (Access Code: `rue4`) **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: ```bash # 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](https://drive.google.com/file/d/1ER_mV2CCsdnF-q3Aoy7loJ2DZXI95M9M/view?usp=drive_link) and place it in the `./checkpoint/` directory. * **Demo Images:** Download the demo images from [Google Drive](https://drive.google.com/file/d/1H7sq-KeK39OfILdoY7MALO6sQqrgaVwf/view?usp=drive_link) 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](https://github.com/MIC-DKFZ/nnUNet), i.e., structured as follows: ```text 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: ```sh # 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]. ```python 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](https://github.com/JJGO/UniverSeg/tree/main) and [Neuralizer](https://github.com/SteffenCzolbe/neuralizer). We extend our gratitude for their significant contributions to the field.