[MedIA 2026] Official PyTorch Implementation of "Test-Time Generative Augmentation for Medical Image Segmentation" (Medical Image Analysis).

• [2025.12.25] 🎄 Merry Christmas! We have updated the Citation and official link. The paper is now fully indexed as Medical Image Analysis (MedIA), Vol. 109, 2026.
• [2025.12.02] 🎉 Accepted! Our paper has been formally accepted by Medical Image Analysis (MedIA)!
• [2025.07.15] 🛠️ Code Release: We have released the complete training codes and inference demos.

Test-Time Generative Augmentation (TTGA) is a novel approach to enhance medical image segmentation during test time. Instead of employing handcrafted transforms or functions on the input test image to create multiple views for test-time augmentation, this approach advocate for the utilization of an advanced domain-fine-tuned generative model, e.g., diffusion models, for test-time augmentation. Hence, by integrating the generative model into test-time augmentation, we can effectively generate multiple views of a given test sample, aligning with the content and appearance characteristics of the sample and the related local data distribution.

The proposed pipeline of three key workflows are presented. The test image is processed through a sequence of steps to generate a noise image at a designated step count. Using this noise image, a one-step denoising process is employed to refine a trainable null-text embedding, enabling the stable generation of results that closely resemble the initial image. In the augmentation generation phase, this null-text embedding, guided by semantic and regional information, is leveraged to produce a series of augmented images.

This repository requires CUDA 12.1, Python 3.9 and Pytorch 2.1. To install the latest version run:

1. Create Conda Environment and Install Dependencies

conda create --name ttga python=3.9
conda activate ttga
pip install -r requirements.txt

2. Install diffusers

• Download diffusers==0.27.0 and extract its contents. You can do this using command-line tools:

wget https://pypi.org/packages/source/d/diffusers/diffusers-0.27.0.tar.gz
tar -xzf diffusers-0.27.0.tar.gz
mv diffusers-0.27.0 diffusers
cd diffusers
python setup.py install
cd ..

This section outlines the steps to run the processes described in this repository, using the polyp dataset as an example.

Step 1: LoRA Fine-tuning of Stable Diffusion First, fine-tune Stable Diffusion using LoRA on the polyp dataset:

bash scripts/lora_polyp.sh

Step 2: One-step Null-Text Optimization Next, perform one-step null-text optimization on test images to extract their identity information.

Step 3: Generative Augmentation Finally, apply generative augmentation to the test images.

Batch Operations and Demos:

• For batch operations of Step 2 and Step 3, please refer to main.py. This script provides a comprehensive implementation for processing multiple images.
• A demonstration showcasing these processes can be found in demo.ipynb. This notebook provides an interactive example of how to apply these techniques.

✨ Optic Disc and Cup Segmentation

✨ Polyp Segmentation

✨ Skin Lesion Segmentation

💕 SOTA segmentation models with codes, datasets and open-source parameters. (Thanks!)

If you find our work useful in your research, please consider citing:

@article{ma2026103902,
  title = {Test-time generative augmentation for medical image segmentation},
  journal = {Medical Image Analysis},
  volume = {109},
  pages = {103902},
  year = {2026},
  issn = {1361-8415},
  doi = {10.1016/j.media.2025.103902},
  author = {Ma, Xiao and Tao, Yuhui and Zhang, Zetian and Zhang, Yuhan and Wang, Xi and Zhang, Sheng and Ji, Zexuan and Zhang, Yizhe and Chen, Qiang and Yang, Guang}
}