Destructive augmentations for prostate lesion segmentation on heterogeneous biparametric MRI data

Introduction

Biparametric MRI (bpMRI) is increasingly used in the management of suspected clinically significant prostate cancer (csPCa). Prostate lesions often appear small and have low contrast with the surrounding tissue, as the signal of malignant tissue overlaps with that of benign tissue. For radiology reporting and MRI-ultrasound fusion biopsy guidance, radiologists contour these lesions manually, a time-consuming process that shows only moderate inter-rater agreement [2]. Moreover, bpMRI protocols vary across hospitals. As a result, publicly available datasets are heterogeneous in both acquisition and annotation, complicating the development of robust deep learning models.

Previous studies highlight the difficulty of prostate lesion segmentation on MRI. The authors of the Prostate158 dataset [2] report baseline Dice scores of 0.450 and 0.400 against two expert readers. Wu et al. [3] outperform the baseline against both readers and report a Dice of 0.513 using arbitrated labels that select the preferred annotation from the two readers. Alzate-Grisales et al. [4] achieve a Dice of 0.467 using a SAM-UNETR model on a combined Prostate158 and PI-CAI cohort. Pooch et al. [5] train an nnU-Net model on a large cohort including PROMIS [6], Prostate158, and PI-CAI, but report only detection metrics. Direct comparison among methods is challenging as evaluation datasets, code, and model weights are often unavailable. To our knowledge, no prior study has trained a segmentation model using all currently available public datasets, including UCL Prostate-MRI-US-Biopsy (PMUB) [7], PROMIS [6], PI-CAI including annotations by Pooch et al. [5], and Prostate158.

Data augmentation is a common strategy for improving robustness to variations in image appearance. Random convolutions [8] perturb image intensities and have been proposed to improve generalization by encouraging shape-based representations, enabling out-of-domain segmentation. Recent work by Molinier et al. [9] similarly improves out-of-domain segmentation performance on a spine segmentation task, without a substantial decrease in in-domain performance through the careful compilation of an extensive list of augmentations, including random convolutions. GIN-IPA also leverages random convolutions and pairs them with a second augmentation aimed at reducing spurious correlations between items in images [10]. These works focus on improving generalization to out-of-domain data, such as CT to MRI generalization or generalization to different MRI sequences. However, imaging data of the same modality can still differ substantially in acquisition, for example in the choice of high b-value for diffusion-weighted imaging (DWI) in prostate MRI, the field strength and the scanner manufacturer.

This work investigates prostate lesion segmentation on bpMRI in multi-dataset training, where cohorts differ in both acquisition and annotation. Given the acquisition differences, we suggest that random convolutions, used for generalization to out-of-domain settings, can be beneficial even in this in-domain task. Additionally, as lesion extent is ambiguous and reader-dependent [2], exact alignment between the image and the ground truth may be less critical. Therefore, we propose the random wavelet augmentation, which perturbs both this image-label alignment and image intensities.

Dice-CE loss is a suboptimal objective for small lesion segmentation, as it is computed over all lesion voxels at once, allowing large lesions to dominate. Instance-level losses address this by considering each lesion individually, e.g., Blob loss [11], Connected-component-Dice-CE (CC-Dice-CE) [12] and Bidirectional CC (BiCC) loss [13]. Complementarily, BackSplit [1] splits the background into different anatomical classes, providing additional training signal that has been shown to improve segmentation of small lesions against large organ backgrounds. Combining instance-level losses with BackSplit remains unexplored.

Our work’s main contributions are:

1. Using 1,579 lesion-containing cases, we train a segmentation model on the largest compiled cohort used for prostate lesion segmentation to date. Our final model achieves state-of-the-art performance on the test set of Prostate158 [2] with a Dice score of 0.614 against the Reader 1 annotations, compared to their 0.450 baseline.

2. Inspired by random convolutions [8] and Random Filter Loss (RaFL) [14], we adapt random wavelets as an augmentation, where the MRI sequences are filtered with a randomly constructed wavelet filter. We show that these destructive augmentations can improve performance when training on heterogeneous labels and acquisition data.

3. We evaluate instance-level losses for prostate lesion segmentation and show that CC-Dice-CE complements BackSplit [1], yielding improved performance over either approach alone.

Methods

We build our pipeline on nnU-Net [15], a self-configuring segmentation framework for medical imaging, and use the recommended ResEnc-M preset [16] as our baseline. We additionally adopt BackSplit [1], adding the prostate mask to the label map, with lesion labels taking precedence. The prostate class serves as an auxiliary training target and is not used at inference. To study augmentation strategies, we compare random wavelets and random convolutions against this nnU-Net-BackSplit baseline (§ 1.1). We then investigate whether instance-level losses can further improve performance with and without BackSplit (§ 1.2).

Random wavelet augmentation

The random wavelet augmentation replaces each sequence with the response of a randomly composed wavelet filter, with

a new filter drawn per sequence. Each sequence is decomposed at four levels $j$ with dilations of 1, 2, 4 and 8, along the three axes $a$ using separable 1D filters with circular padding. At each level, we combine the three single-axis detail bands $D_{j}^{(a)}$ along a random unit direction $u_j \sim\mathcal{U}(S^2)$, discarding the mixed-axis detail bands and low-pass residual. The level responses are then summed with random weights $v_j \sim\mathcal{N}(0,1)$, $F = \sum_{j=1}^{4} v_j \sum_a u_{j,a}D_{j}^{(a)}$. The result is min-max rescaled to the intensity range of the original sequence. We evaluate four wavelet types: db1 (Haar, 2 taps), db2 (4 taps), db4 (8 taps, smooth but strongly asymmetric) and Symlet sym4, a near-symmetric counterpart of db4. Figure 1 shows an example of a sym4 augmentation, illustrating that the filter not only distorts textures but also displaces edges.

We compare these to the default and DA5 augmentations included in nnU-Net. The default pipeline consists of standard medical imaging augmentations, including mirroring, rotation, brightness and contrast changes, Gaussian noise and blurring. DA5 increases the strength and/or probability of the transformations and adds transformations such as blank rectangles or sharpening.

On top of the nnU-Net default augmentations, we add either random wavelets or random convolutions with a probability of 0.5. For random convolutions, we test two configurations: 3 layers (l=3) of 3 × 3 × 3 (k=3) convolutions and 5 layers (l=5) of 5 × 5 × 5 (k=5).

Instance-level losses

To better account for lesion-level imbalance, we investigate augmenting the Dice and Cross-Entropy loss (Dice-CE) with instance-level losses. As instance-level losses do not benefit large components in the BackSplit class, we apply the Dice-CE loss to all classes, but limit the instance-level losses to the lesions only, as follows:

where $Y$ is the ground-truth label map for all classes and $\hat{Y}$ is the network’s softmax output, $y_{\ell}$ is the binary reference lesion mask and $\hat{y}_{\ell}$ is the predicted lesion probability, i.e. the lesion channel of $\hat{Y}$, with the prostate and remaining background channel merged into a single class. $\mathcal{L}_{CE}$ and $\mathcal{L}_{SD}$ are nnU-Net’s cross-entropy and soft Dice losses. $\mathcal{L}_{inst}$ is one of the three instance-level losses described below. We use nnU-Net’s deep supervision implementation and weights for every loss term.

We implement three instance-level losses: Blob loss [11] with Dice-CE, CC-Dice-CE [12] and BiCC-Dice-CE [13]. Blob loss iterates over all ground-truth lesions, masks out the other ground-truth lesions and calculates the loss between that lesion and the predicted mask. The results are averaged over all lesions. Alternatively, CC-Dice-CE creates Voronoi cells for each ground-truth label and calculates the Dice-CE loss within each cell. Bidirectional CC-Dice-CE (BiCC-Dice-CE) additionally divides the predicted mask into Voronoi cells and matches them to the ground-truth, reducing false positives.

Experimental setup

Data

We use five prostate bpMRI datasets: a dataset derived from PROMIS [6], PI-CAI [17], Prostate-MRI-US-Biopsy (PMUB) [7], Prostate158 [2] and an internal cohort, described in Table 2. For Prostate158, the official test split is used. For the other datasets, we define the split ourselves and make them available as a reproducible evaluation. All datasets include three sequences: T2W, ADC and high-b-value DWI. ADC and DWI sequences are resampled to the T2W grid. For PI-CAI, we use the AI-generated prostate masks [18], the lesion segmentations included in the dataset and additional annotations by Pooch et al. [5]. For Prostate158, we use the annotations from Reader 1 [2]. This dataset only annotates

Dataset	Train/Val/Test	Field (T)	High $b$ (s/mm$^2$)	Lesion vol Median (mL)
Internal	428 / 52 / 117	3	1400	0.182
PI-CAI [17]	296 / 43 / 86	1.5/3	$\ge1000$	1.055
PROMIS [6]	317 / 33 / 87	1.5	1400	1.121
P158 [2]	70 / 12 / 19	3	1000	0.441
PMUB [7]	468 / 67 / 133	3	1400	0.347
Total	1579 / 207 / 442

lesions of PI-RADS $\ge4$. For our internal data we use all lesions annotated by clinicians during interpretation of the case, which generally includes only lesions with PI-RADS $\ge3$. In PMUB, all annotated lesions correspond to lesions targeted for biopsy, while in PI-CAI, lesions with ISUP grade $\ge2$ were annotated. As illustrated in Table 2, lesion sizes vary substantially between datasets, which can be attributed to different labeling protocols. Cases without any annotated lesions are excluded.

Training and evaluation

We follow the nnU-Net training recipe, using 1000 epochs of 250 steps at batch size 2. All configurations are trained on the full training set described in Table 2 for a single fold. nnU-Net’s automatic configuration yields a patch size of 224 × 224 × 32 voxels at a spacing of 0.5 × 0.5 × 1.5mm.

We evaluate all models with the final training checkpoint on the same held-out test sets (see Table 2). We report the Dice score averaged across test cases and precision, recall and F1 score calculated across all individual lesions. A predicted lesion is counted as a true positive if it overlaps a ground-truth lesion with a Dice score above 0.1.

Results

Combined Improvements and Overall Performance: We first evaluate the cumulative effect of our proposed modifications on top of the nnU-Net baseline. Table 1 shows the additive effect of consecutively adding BackSplit and our two technical contributions. BackSplit [1] improves Dice performance by +.011 and the F1 score by +.070 over plain nnU-Net. CC-Dice-CE further improves this result by +.006 Dice, and the sym4 wavelet by another +.014. Both additions increase recall, but lower precision, reducing the F1 score from .596 for the BackSplit-only model to .571 for our final model. For our task, we prioritize recall over precision, as false negatives have a higher impact than false positives. Compared to the nnU-Net baseline, our final model still sees an improvement in Dice of +0.031 and F1 of +0.045. On the Prostate158 dataset, our final model achieves a Dice of 0.614 against the labels of Reader 1, compared to the baseline of 0.450 provided by Adams et al. [2] and 0.506 by Wu et al. [3], setting a new state-of-the-art on this dataset.

All

Internal

PROMIS

PI-CAI

Prostate158

PMUB

Dice

Detection (R/P/F1)

Dice

Detection (R/P/F1)

Dice

Detection (R/P/F1)

Dice

Detection (R/P/F1)

Dice

Detection (R/P/F1)

Dice

Detection (R/P/F1)

nnU-Net ResEnc-M

.391

.624/.454/.526

.325

.519/.392/.447

.341

.560/.398/.465

.497

.821/.446/.578

.609

.870/.800/.833

.383

.651/.552/.597

+ BackSplit gland

.402

.594/.597/.596

.336

.481/.517/.499

.351

.520/.523/.522

.533

.811/.706/.755

.603

.826/.864/.844

.380

.628/.639/.633

+ CC

.408

.635/.537/.582

.339

.524/.439/.478

.364

.593/.511/.549

.530

.842/.645/.731

.614

.913/.808/.857

.388

.640/.567/.601

+ random sym4 wavelet

.422

.682/.492/.571

.346

.535/.388/.450

.397

.687/.510/.585

.532

.863/.599/.707

.614

.913/.808/.857

.405

.703/.492/.579

Effect of Destructive Augmentations: Table 3 compares the different augmentation strategies applied to our nnU-Net baseline optimized with BackSplit. These results show a trend that lesion segmentation performance improves as the transformations become more aggressive. While DA5 augmentations provide a substantial gain over the standard nnU-Net pipeline (+0.011 Dice), even larger improvements are achieved with random convolutions (+0.013 Dice for $k=5$, $l=5$) and long-range wavelet perturbations. Among the wavelet-based methods, Dice performance increases progressively from Haar to db2, db4, and sym4, with sym4 yielding the largest overall improvement (+0.016 Dice). The benefits of augmentation are most pronounced for PROMIS, the only cohort acquired entirely at 1.5T, suggesting improved robustness to underrepresented lower field strength images in our training data.

Augmentation family	All Dice	All Detection (R/P/F1)	Internal Dice	Internal Detection (R/P/F1)	PROMIS Dice	PROMIS Detection (R/P/F1)	PI-CAI Dice	PI-CAI Detection (R/P/F1)	Prostate158 Dice	Prostate158 Detection (R/P/F1)	PMUB Dice	PMUB Detection (R/P/F1)
no random filter (nnU-Net default DA)	.402	.594/.597/.596	.336	.481/.517/.499	.351	.520/.523/.522	.533	.811/.706/.755	.603	.826/.864/.844	.380	.628/.639/.633
DA5	.413	.624/.533/.575	.339	.508/.431/.467	.375	.587/.471/.522	.536	.842/.684/.755	.603	.870/.800/.833	.395	.628/.584/.605
random conv $k = 3$, $l = 3$	.407	.622/.494/.551	.329	.492/.392/.436	.384	.613/.511/.558	.526	.832/.590/.690	.597	.826/.760/.792	.385	.628/.498/.555
random conv $k = 5$, $l = 5$	.415	.658/.521/.581	.340	.514/.438/.473	.385	.653/.513/.575	.532	.863/.636/.732	.618	.913/.778/.840	.397	.669/.511/.579
random wavelet, Haar (db1)	.411	.621/.534/.574	.339	.470/.424/.446	.383	.640/.519/.573	.539	.842/.667/.744	.624	.870/.870/.870	.378	.610/.541/.574
random wavelet, db2	.415	.646/.531/.583	.354	.535/.446/.486	.380	.633/.503/.560	.545	.863/.678/.759	.632	.913/.840/.875	.378	.622/.525/.569
random wavelet, db4	.417	.645/.523/.577	.345	.508/.437/.470	.383	.613/.503/.553	.547	.874/.680/.765	.619	.913/.778/.840	.390	.657/.504/.571
random wavelet, sym4	.418	.648/.536/.587	.356	.519/.442/.478	.386	.627/.516/.566	.542	.853/.659/.743	.615	.913/.808/.857	.385	.657/.543/.595

In contrast to the improvements in Dice, the F1 score decreases with the more aggressive DA5, random convolution and wavelet augmentations. This reduction is generally driven by lower precision despite higher recall.

For sym4 wavelets, however, the decrease is small compared to the F1 improvement obtained from BackSplit.

These findings suggest that even highly destructive augmentations can remain beneficial in the presence of noisy labels and heterogeneous datasets. One possible explanation is that such augmentations act as a regularizer, reducing overfitting to specific image characteristics, similar to dropout in neural networks.

Effect of Instance-Level Losses: We evaluate three instance-level losses on top of Dice-CE, both with and without BackSplit. Table 4 shows that CC-Dice-CE improves Dice performance over the Dice-CE baseline and achieves the best overall performance when combined with BackSplit. In contrast, Blob-loss and BiCC do not provide consistent improvements across datasets.

Loss	All	Internal	PROMIS	PI-CAI	P-158	PMUB
lesion vs background
Dice-CE	.391	.325	.341	.497	.609	.383
+ blob	.392	.319	.336	.503	.615	.390
+ CC	.398	.328	.354	.500	.615	.391
+ BiCC	.388	.319	.326	.515	.591	.379
BackSplit gland
Dice-CE	.402	.336	.351	.533	.603	.380
+ blob	.401	.331	.351	.524	.611	.385
+ CC	.408	.339	.364	.530	.614	.388
+ BiCC	.397	.334	.345	.514	.624	.379

Conclusion

In this work, we compile a large prostate lesion segmentation cohort spanning five heterogeneous datasets with substantial differences in acquisition protocols, lesion characteristics, and annotation practices. Despite this challenging setting, our final model achieves state-of-the-art performance on the Prostate158 test set, reaching a Dice score of 0.614. We find that destructive wavelet-based augmentations improve robustness to dataset heterogeneity, while CC-Dice-CE provides additional gains beyond standard Dice-CE optimization and complements BackSplit. Further investigations could study the effect of perturbing other wavelet sub-bands, different combinations of sub-bands and different wavelet types. Additionally, evaluations on different tasks and in a true out-of-domain setting could further establish the value of the random wavelet augmentations.

Compliance with ethical standards

For the use of our internal data, we obtained approval of Weill-Cornell Medicine and New-York Presbyterian hospital review board. All other data are publicly accessible.

Acknowledgments

This work is supported by the National Institute of Health grants K25CA283145, R01AG053949, R01AG064027, and R01AG070988. Thibault Heintz is supported by the Belgian American Educational Foundation. Daniel Margolis reports clinical advisory to Stratagen Bio/Ravada Health, holding stock options. Mert R. Sabuncu is an advisor to Cleerly Inc. and Medara Health and holds equity interests in both. We use the AI-tool Claude for AI assisted programming and grammar checking.

References