WAAU-Net

Wavelet-domain adversarial adaptation for cross-cohort generalization in breast ultrasound tumor segmentation.

The gist

Before building anything, this project asks a diagnostic question first. When a segmentation model trained on one hospital's scans fails on another's, where does that failure actually live? A radiomics analysis went looking for the answer. It checked 650 image features across three hospital cohorts, hunting for which ones differed significantly between them. The divergence was not scattered evenly across the image. It concentrated in one place. Sixty percent of the features that differed significantly between cohorts came from wavelet-decomposed versions of the image, not from the raw pixels.

So WAAU-Net does not try to scrub domain information out of the whole image at once, the way most adversarial domain-adaptation methods do. It decomposes the image into frequency bands first. Then it normalises only the bands the data actually pointed to.

Research question

Breast-ultrasound segmentation models are usually trained and tested within one hospital cohort. Move that same model to a different hospital and its accuracy drops, sometimes badly. Cohorts differ in scanner, in acquisition protocol, and in patient population all at once, and that shift tends to hit underrepresented groups hardest.

Existing domain-adaptation methods try to fix this in the spatial domain. They suppress domain information as one undifferentiated whole. There is no way to tell, with that approach, whether diagnostically useful structure got suppressed right along with the noise. The real question this project asks is narrower and more answerable: instead of guessing where domain-specific noise lives in an image, can you find out directly, and act only on that part?

Solution

Each image first passes through a discrete wavelet transform, which splits it into four subbands: LL, LH, HL, and HH. An Adaptive Wavelet Normalization module then learns a gate for each subband. That gate blends the raw signal with an instance-normalised version of it. A high gate value keeps the raw signal mostly intact. A low gate value replaces it with the normalised form instead.

The whole network, an Attention U-Net backbone reading these gated subbands, is trained jointly with a second network, a discriminator that tries to guess which hospital an image came from. The main encoder is pushed to make that guess fail. The intuition behind the design mirrors the radiomics finding directly. Tumour localisation mostly lives in the LL and LH bands, which capture intensity contrast and true tissue edges. The HL and HH bands carry more of the cohort-varying texture and acquisition artifact. Selectively normalising only the high-frequency bands should strip out cohort-specific noise without erasing the diagnostic signal underneath it.

From radiomics finding to framework to clinical outcome

Stage 1 is the diagnostic radiomics analysis. Stage 2 is WAAU-Net itself: wavelet decomposition feeding Adaptive Wavelet Normalization, which keeps LL and LH at high attention and suppresses HL and HH, jointly trained against a gradient-reversed discriminator. Stage 3 is what that buys clinically.

Results

Three cohorts were used, restricted to malignant cases and balanced to 98 images each: AIIMS New Delhi, the open-access BUSI cohort from Egypt, and the open-access BrEaST cohort from Poland. Internal performance was measured with 5-fold cross-validation. External performance used a leave-one-cohort-out setup, so the test cohort was never seen during training at all.

Internally, WAAU-Net is competitive with the strongest baseline, scoring a Dice of 0.73 against nnU-Net's 0.74. The real difference shows up externally, on the held-out BrEaST cohort, where WAAU-Net holds the smallest internal-to-external gap of any model tested, at 0.090.

Predictions and GradCAM activations across four cases

WAAU-Net's activation consistently converges on one compact, anatomically plausible region, where several baselines spread activation across the image or split it across multiple areas.

An ablation shows where that advantage actually comes from. Most of it traces back to the adversarial training itself, not to the wavelet targeting on its own.

And here is the payoff for the diagnostic detour this page started with. The AWN module's learned gate values stayed stable across every configuration tested, and they land almost exactly where the radiomics analysis predicted they should.

What next

The real value of this result is that it turns a guess into a targeted fix. Most domain-adaptation work suppresses domain information across an entire image and hopes the diagnostic signal survives. WAAU-Net instead starts from a measurement of exactly where cohort differences live, then builds a mechanism aimed only at that location. The fact that the network's own learned gates converge on almost the same split the radiomics analysis found beforehand is a strong sign that the targeting is real, not a coincidence of training.

That kind of targeted normalisation should generalise past breast ultrasound. Any imaging setting where a model has to move between scanners, protocols, or populations faces the same underlying problem: a mix of genuinely useful structure and cohort-specific noise sitting in the same pixels. A frequency-domain measurement step before adaptation is a pattern that could transfer wherever that mix shows up.

The honest limitation is scale. These cohorts are small, malignant-only, and the study cannot separate the effects of scanner, population, and acquisition protocol from one another, since all three vary together across cohorts. Broader validation, with more cohorts and enough cases to test benign and malignant findings separately, is the identified next step.

Status & links

Manuscript under review at Biomedical Signal Processing and Control. Read the manuscript. Code is available at github.com/ViciousDove19/WAAUNet.