Test-time generative augmentation for medical image segmentation

TL;DR

This paper introduces Test-Time Generative Augmentation (TTGA), a novel method using domain-fine-tuned generative models to improve medical image segmentation accuracy and uncertainty estimation during inference.

Contribution

We propose TTGA, a new test-time augmentation strategy leveraging diffusion models for region-specific, contextually relevant augmentations in medical image segmentation.

Findings

Improves segmentation accuracy with DSC gains up to 2.3%.

Enhances pixel-wise error estimation with DSC gains up to 29%.

Demonstrates effectiveness across multiple datasets and segmentation tasks.

Abstract

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.