++nnU-Net: Scaling nnU-Net with Prefix-Based Data Augmentation

2026-06-09

AI Summary

This work addresses the challenge of limited annotated data in medical image segmentation, exacerbated by privacy constraints, by proposing a scalable prefix-based data augmentation strategy that enhances few-shot 2D segmentation performance without altering the nnU-Net architecture. The approach leverages a two-stage image registration pipeline to generate deformed images along with their corresponding labels, and further incorporates synthetic mask generation and an efficient disk management mechanism to substantially increase training data diversity. Experimental results across five 2D medical image datasets demonstrate significant improvements over the original nnU-Net baseline, with Dice similarity coefficients increasing by up to approximately 22%.

Abstract

The nnU-Net has demonstrated continuous success in medical segmentation tasks, which heavily rely on the availability and diversity of annotated biomedical data. However, assembling medical imaging cohorts remains challenging due to numerous factors such as privacy regulations and annotation costs. As a result, data augmentation plays a crucial role in increasing data availability while maintaining anatomical feasibility. Hence, we propose the ++nnU-Net, a novel data augmentation module based on image registration that operates prior to preprocessing and training take place. Our framework was evaluated across five different 2D datasets. In this workflow, image data go through a two-stage registration process, generating new warped images. The transformations are then applied to the respective segmentation. In addition, the pipeline computes available disk space, generates supplementary binary synthetic masks and generates checkpoints. We demonstrate that the ++nnU-Net outperforms the nnU-Net baseline, yielding improvements in Dice Similarity Coefficient scores. In the most prominent cases, we observe performance gains of approximately 22\%. These findings highlight the effectiveness of registration-based data augmentation, particularly for 2D medical imaging datasets and suggest that the ++nnU-Net provides a practical and scalable approach for enhancing segmentation performance in data-limited settings. The source code for the ++nnU-Net is available at: https://github.com/sofia-adelie/plusplusnnunet.git

Problem

• medical image segmentation
• data scarcity
• data augmentation
• annotated biomedical data
• limited datasets

Innovation

• prefix-based data augmentation
• image registration
• medical image segmentation
• nnU-Net
• synthetic mask generation