💻 computer science

Decoupling Wavelet Sub-bands for Single Source Domain Generalization in Fundus Image Segmentation

This paper introduces WaveSDG, a novel wavelet-guided network for single-source domain generalization in fundus image segmentation that decouples anatomical structure from domain-specific appearance via a specialized WISER module, achieving superior accuracy and robustness across multiple unseen target datasets compared to state-of-the-art methods.

Original authors: Shramana Dey, Varun Ajith, Abhirup Banerjee, Sushmita Mitra

Published 2026-04-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shramana Dey, Varun Ajith, Abhirup Banerjee, Sushmita Mitra

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to identify the "cup" inside a "saucer" in a photograph of the back of a human eye (the retina). This is crucial for doctors to spot diseases like glaucoma.

The problem is that the robot learns on photos taken with one specific camera in one specific hospital. When you show it photos from a different hospital, taken with a different camera, under different lights, or by a different doctor, the robot gets confused. It starts seeing the style of the photo (the lighting, the color tint, the graininess) as part of the structure of the eye. It might think a dark shadow is a disease, or miss a boundary because the colors look different.

This paper introduces a new method called WaveSDG to fix this. Here is how it works, using simple analogies:

1. The Problem: The "Style" vs. "Structure" Mix-Up

Think of a fundus image like a hand-drawn map.

Structure (Anatomy): The actual roads, rivers, and landmarks (the optic disc and cup). This is what matters for the diagnosis.
Style (Appearance): The color of the ink, the type of paper, or the lighting in the room where the map was drawn.

Current AI models often get tricked by the "style." If they train on maps drawn in blue ink, they might fail when shown a map drawn in red ink, even if the roads are in the exact same place.

2. The Solution: The "Wavelet" Magic Filter

The authors use a mathematical tool called Wavelet Decomposition. Imagine taking that hand-drawn map and running it through a special machine that splits the image into four different "layers" or "sub-bands":

The "LL" Layer (Low-Low): This is the blurred, smooth version of the map. It keeps the big picture and the general shape of the roads (the anatomy) but loses the sharp details.
The "LH" and "HL" Layers: These are the edge detectors. They highlight the horizontal and vertical lines—the sharp boundaries of the roads and rivers.
The "HH" Layer: This is the static and noise. It's like the grainy dust on the paper or the random scratches. It usually doesn't contain useful map information.

3. The "WISER" Module: The Smart Editor

The paper introduces a new module called WISER (Wavelet-based Invariant Structure Extraction and Refinement). Think of WISER as a very smart editor who looks at these four layers and decides what to keep and what to throw away before the AI makes its final decision.

Cleaning the "LL" (The Big Picture):
The editor looks at the smooth "LL" layer. It knows this layer holds the shape of the eye, but it also holds the "style" (like the specific lighting of the hospital). WISER splits this layer into two: one part that keeps the shape (anatomy) and one part that captures the lighting (style). It then throws away the lighting part, keeping only the pure shape.
- Analogy: It's like taking a photo of a building, removing the sunset glow and the shadows, and keeping only the outline of the building.
Polishing the "LH" and "HL" (The Edges):
The editor looks at the edge layers. Sometimes, bad lighting makes the edges look weak or broken. WISER has two tools:
1. The Edge Booster: If an edge is faint (like a faint road line), it turns up the volume to make it visible.
2. The Edge Selector: If an edge is just random noise (like a scratch on the lens), it silences it. It only keeps the edges that match the "shape" it found in the previous step.
- Analogy: It's like a sound engineer turning up the volume on a singer's voice (the real edge) while turning down the background noise (the fake edge).
Deleting the "HH" (The Noise):
The editor simply throws away the "HH" layer entirely because it's mostly just noise and artifacts.

4. The Result: A Robust AI

After WISER cleans and refines these layers, it stitches them back together. The AI now sees the eye not as "a photo taken with a Canon camera in a dark room," but as "a set of pure anatomical shapes and clear boundaries."

Because the AI has learned to ignore the "style" (the camera, the lighting, the hospital) and focus only on the "structure" (the actual eye parts), it works perfectly well even when shown photos from completely new hospitals or cameras it has never seen before.

What the Paper Found

The researchers tested this on a "source" dataset (one hospital) and then challenged it with five completely different "target" datasets (other hospitals with different cameras).

The Winner: WaveSDG (their new method) consistently beat seven other top-tier methods.
The Proof: It didn't just guess the right area; it drew the boundaries with much higher precision and stability. It was less likely to be confused by weird lighting or different camera types.
Efficiency: The paper notes that this "smart editor" is lightweight. It doesn't require a supercomputer to run; it adds very little extra work to the system, making it practical for real-world use.

In short, the paper teaches the AI to look past the "clothing" (style) of the image to see the "body" (anatomy) underneath, ensuring it can recognize the eye correctly no matter where the photo was taken.

1. Problem Statement

Domain Generalization (DG) in Fundus Imaging:
Deep learning models for retinal fundus image segmentation (specifically for the Optic Disc and Optic Cup) suffer from performance degradation when deployed on data from unseen domains (different hospitals, devices, or lighting conditions). This is caused by domain shift, where variations in acquisition conditions (illumination, sensor noise, texture) create a distribution mismatch between training (source) and testing (target) data.

Limitations of Existing Approaches:

Domain Adaptation (DA) & Test-Time Adaptation (TTA): Require access to target domain data during training or inference, which is often restricted by privacy regulations and annotation costs.
Single-Source Domain Generalization (SDG): The most realistic scenario where models are trained on one source dataset and must generalize to unseen targets. Existing SDG methods often fail because:
- Data Augmentation: Synthetic transformations often corrupt critical anatomical structures.
- Representation Learning: Disentangling "style" (appearance) from "content" (anatomy) is difficult as they are inherently coupled in medical images.
- Frequency Methods: Fourier-based methods lack spatial locality, making them poor at modeling the spatial organization of anatomical structures.

2. Methodology: WaveSDG and WISER

The authors propose WaveSDG, a segmentation network built on a U-shaped encoder-decoder backbone (ResNet18 encoder) that integrates a novel module called Wavelet-based Invariant Structure Extraction and Refinement (WISER).

Core Concept: Wavelet Sub-band Decoupling

Instead of treating all image features equally, the WISER module applies a Depthwise Haar Discrete Wavelet Transform (DWT) to the encoder feature maps. This decomposes features into four sub-bands, each with distinct semantic roles:

LL (Low-Low): Captures global anatomical structure and low-frequency content.
LH/HL (Low-High / High-Low): Capture horizontal and vertical edge responses (boundaries).
HH (High-High): Dominated by sensor noise and artifacts.

The WISER Module Architecture

The WISER module processes encoder features layer-by-layer before they are passed to the decoder via skip connections. It performs three specific operations:

Anatomy-Style Decoupling (LL Sub-band):
- The LL sub-band is decomposed into Style ( $f^s$ ) and Anatomy/Content ( $f^c$ ) representations using identical convolutional layers and Instance Normalization (IN).
- An Orthogonal Loss is applied to enforce decorrelation between $f^s$ and $f^c$ , ensuring the anatomy representation is independent of domain-specific appearance (e.g., lighting, color tone).
- The refined LL sub-band is reconstructed using the normalized content representation.
Selective Edge Refinement (LH and HL Sub-bands):
- Edge Booster (EB): Computes normalized edge energy to adaptively amplify weak edge responses caused by low contrast or poor image quality.
- Edge Selector (ES): Uses the anatomy-aware content feature ( $f^c$ ) to generate a spatial gate. This gate selectively retains edges aligned with anatomical structures while suppressing background noise and spurious edges.
- The LH and HL sub-bands are refined by multiplying them with the outputs of both the EB and ES.
Noise Suppression (HH Sub-band):
- The HH sub-band, which is dominated by noise, is completely suppressed (set to zero) to prevent noise propagation.
Reconstruction:
- The filtered sub-bands are reconstructed via Inverse DWT (IDWT).
- A residual connection combines the filtered features with the original features to ensure training stability.

Training Strategy

Loss Function: The model uses a Multi-class Dice Loss for segmentation, augmented with Deep Supervision (DS) to enforce boundary consistency across scales.
Orthogonality Constraint: A linear ramp schedule is used to gradually introduce the orthogonal loss, allowing the model to first learn task-specific representations before enforcing strict decoupling.

3. Key Contributions

Novel WISER Module: A lightweight, resource-efficient module that decouples anatomical structure from domain-specific appearance using wavelet sub-band decomposition.
Semantic Sub-band Processing: Unlike previous wavelet methods that treat all sub-bands equally, WISER explicitly leverages the semantic differences: refining LL for global anatomy, enhancing LH/HL for boundaries, and suppressing HH for noise.
Edge Booster & Selector: A dual mechanism that adaptively amplifies weak edges and filters out non-anatomical noise, ensuring precise Region of Interest (ROI) delineation even under large domain shifts.
WaveSDG Architecture: A complete SDG framework that prevents the propagation of domain-specific appearance cues through skip connections, a common failure point in standard U-Nets.

4. Experimental Results

Datasets:

Source: Refuge dataset (840 training, 120 validation, 240 test).
Target (Unseen): 5 datasets including Drishti-GS, Gamma, and the multi-domain Cháks.u (Bosch, Forus, Remidio), covering various camera types (clinical and handheld).

Performance Metrics:

Dice Score Coefficient (DSC): Measures overlap.
95th Percentile Hausdorff Distance (HD95): Measures boundary alignment precision.

Key Findings:

Superior Accuracy: WaveSDG consistently outperformed seven State-of-the-Art (SOTA) methods (including BigAug, C2SDG, FreeSDG, and EGSDG) across all target datasets.
Robustness: It achieved the best balanced Dice scores for both Optic Disc (OD) and Optic Cup (OC) segmentation.
Boundary Precision: WaveSDG achieved the lowest HD95 with reduced variance, indicating superior boundary alignment and stability.
Generalization: The model performed exceptionally well on handheld camera datasets (e.g., Forus, Remidio), which typically suffer from severe domain shifts.
Ablation Studies: Confirmed that the WISER module significantly improves DSC and reduces HD95. The addition of Deep Supervision further enhanced performance.
Efficiency: The module added only 3.24M parameters and a 13.27% computational overhead, maintaining a low inference latency of 14.8ms.

5. Significance

This research addresses a critical bottleneck in medical AI: the inability of models to generalize across different clinical settings without retraining. By exploiting the Discrete Wavelet Transform (DWT), the authors provide a principled way to separate "what" the anatomy is from "how" it looks (style).

Clinical Impact: The method enables the deployment of robust segmentation tools for glaucoma and other retinal diseases using data from a single source, reducing the need for expensive, privacy-sensitive multi-center data collection.
Technical Innovation: It moves beyond simple augmentation or adversarial training, offering a structural solution that respects the physical and semantic properties of retinal images (global layout vs. local edges).
Future Directions: The modular design suggests potential for application in cross-modality tasks (e.g., MRI to CT) or cross-sequence generalization.