Destructive augmentations for prostate lesion segmentation on heterogeneous biparametric MRI data

Abstract

Training prostate lesion segmentation models on biparametric MRI (bpMRI) is limited by small datasets with heterogeneous imaging acquisition protocols and annotations. In this work, we train a segmentation model on a combined cohort of 1,579 lesion-containing training cases from five datasets, representing, to our knowledge, the largest collection of prostate lesion annotations used for this task. Additionally, we investigate two complementary improvements to the widely used nnU-Net framework. First, we introduce a random wavelet augmentation and show that random sym4 (Symlet) wavelets outperform random convolution augmentation in terms of average Dice score, with only a modest reduction in F1 score due to lower precision. Second, we evaluate instance-level losses and show that Connected-Component-based Dice (CC-Dice)-CE loss complements BackSplit, yielding better performance than BackSplit alone. Together, these improvements achieve the best overall performance, reaching a Dice score of 0.614 on the public Prostate158 test set and establishing a new state-of-the-art on this benchmark. We publicly release our training code and model weights.