# Test-time generative augmentation for medical image segmentation

arXiv: https://arxiv.org/abs/2406.17608 · HTML: https://mlsanity.com/paper/2406.17608/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 45 references that are in the map.
- First posted on arXiv 25 Jun 2024.

## Builds on (most cited first, 30 of 33)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Prompt-to-Prompt Image Editing with Cross Attention Control](https://mlsanity.com/paper/2208.01626/) (2022)
- [Null-text Inversion for Editing Real Images using Guided Diffusion Models](https://mlsanity.com/paper/2211.09794/) (2022)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [Diffusion Model-Based Image Editing: A Survey](https://mlsanity.com/paper/2402.17525/) (2024)
- [Unsupervised Medical Image Translation with Adversarial Diffusion Models](https://mlsanity.com/paper/2207.08208/) (2022)
- [Confidence Calibration and Predictive Uncertainty Estimation for Deep Medical Image Segmentation](https://mlsanity.com/paper/1911.13273/) (2019)
- [MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation](https://mlsanity.com/paper/2303.09975/) (2023)
- [DS-TransUNet:Dual Swin Transformer U-Net for Medical Image Segmentation](https://mlsanity.com/paper/2106.06716/) (2021)
- [Prompt Tuning Inversion for Text-Driven Image Editing Using Diffusion Models](https://mlsanity.com/paper/2305.04441/) (2023)
- [Learning with Limited Annotations: A Survey on Deep Semi-Supervised Learning for Medical Image Segmentation](https://mlsanity.com/paper/2207.14191/) (2022)
- [Assessing Reliability and Challenges of Uncertainty Estimations for Medical Image Segmentation](https://mlsanity.com/paper/1907.03338/) (2019)
- [CAMEL: A Weakly Supervised Learning Framework for Histopathology Image Segmentation](https://mlsanity.com/paper/1908.10555/) (2019)
- [Medical Image Segmentation Using Squeeze-and-Expansion Transformers](https://mlsanity.com/paper/2105.09511/) (2021)
- [Label-Free Liver Tumor Segmentation](https://mlsanity.com/paper/2303.14869/) (2023)
- [Source-Free Domain Adaptation for Image Segmentation](https://mlsanity.com/paper/2108.03152/) (2021)
- [Efficient Medical Image Segmentation Based on Knowledge Distillation](https://mlsanity.com/paper/2108.09987/) (2021)
- [Automatic Brain Tumor Segmentation using Convolutional Neural Networks with Test-Time Augmentation](https://mlsanity.com/paper/1810.07884/) (2018)
- [H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation](https://mlsanity.com/paper/2403.13642/) (2024)
- [Contextual Attention Network: Transformer Meets U-Net](https://mlsanity.com/paper/2203.01932/) (2022)

## Built on by (most cited first, 6 of 6)

- [Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoning](https://mlsanity.com/paper/2506.04611/) (2025)
- [A Closer Look at Edema Area Segmentation in SD-OCT Images Using Adversarial Framework](https://mlsanity.com/paper/2508.18790/) (2025)
- [Value bounds and Convergence Analysis for Averages of LRP attributions](https://mlsanity.com/paper/2509.08963/) (2025)
- [Improving MLLM Historical Record Extraction with Test-Time Image](https://mlsanity.com/paper/2509.09722/) (2025)
- [SegTTA: Training-Free Test-Time Augmentation for Zero-Shot Medical Imaging Segmentation](https://mlsanity.com/paper/2604.17451/) (2026)
- [R2AoP: Reliable and Robust Angle of Progression Estimation from Intrapartum Ultrasound](https://mlsanity.com/paper/2605.21099/) (2026)

## Abstract (arXiv)

Medical image segmentation is critical for clinical diagnosis, treatment planning, and monitoring, yet segmentation models often struggle with uncertainties stemming from occlusions, ambiguous boundaries, and variations in imaging devices. Traditional test-time augmentation (TTA) techniques typically rely on predefined geometric and photometric transformations, limiting their adaptability and effectiveness in complex medical scenarios. In this study, we introduced Test-Time Generative Augmentation (TTGA), a novel augmentation strategy specifically tailored for medical image segmentation at inference time. Different from conventional augmentation strategies that suffer from excessive randomness or limited flexibility, TTGA leverages a domain-fine-tuned generative model to produce contextually relevant and diverse augmentations tailored to the characteristics of each test image. Built upon diffusion model inversion, a masked null-text inversion method is proposed to enable region-specific augmentations during sampling. Furthermore, a dual denoising pathway is designed to balance precise identity preservation with controlled variability. We demonstrate the efficacy of our TTGA through extensive experiments across three distinct segmentation tasks spanning nine datasets. Our results consistently demonstrate that TTGA not only improves segmentation accuracy (with DSC gains ranging from 0.1% to 2.3% over the baseline) but also offers pixel-wise error estimation (with DSC gains ranging from 1.1% to 29.0% over the baseline). The source code and demonstration are available at: https://github.com/maxiao0234/TTGA.
