WaveRNet: Wavelet-Guided Retinal Segmentation

• WaveRNet is a wavelet-guided frequency learning framework that decomposes retinal images into low- and high-frequency components for robust, domain-generalized segmentation.
• It leverages specialized modules—including the Spectral-guided Domain Modulator, Frequency-Adaptive Domain Fusion, and Hierarchical Mask-Prompt Refiner—to dynamically fuse domain-specific features and refine segmentation masks.
• Empirical evaluations demonstrate superior cross-domain performance with higher Dice scores compared to SAM-based and U-Net approaches under Leave-One-Domain-Out protocols.

WaveRNet is a wavelet-guided frequency learning architecture designed for multi-source domain-generalized retinal vessel segmentation. Developed to address the pervasive challenge of domain shift due to non-uniform illumination, contrast variation, and the need to preserve fine vessel structures, WaveRNet advances over existing Segment Anything Model (SAM)-based approaches by incorporating explicit frequency-domain decomposition with adaptive domain fusion and hierarchical mask refinement. Its architecture integrates a Spectral-guided Domain Modulator (SDM), Frequency-Adaptive Domain Fusion (FADF), and a Hierarchical Mask-Prompt Refiner (HMPR) for robust generalization across heterogeneous retinal imaging domains (Wang et al., 9 Jan 2026).

1. Core Architectural Components

WaveRNet processes a retinal image $x \in \mathbb{R}^{H \times W \times 3}$ using a series of specialized modules:

• SAM Image Encoder with Adapters: Uses the ViT-B backbone from SAM with lightweight adapters inserted into each transformer block. All SAM weights are frozen. The encoder generates a feature map $\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$, with $H' = H/16$, $W' = W/16$.
• Spectral-guided Domain Modulator (SDM): Decomposes $\mathbf{F}$ into low- and high-frequency branches using learnable convolutional "wavelet" layers. These branches are fused and modulated by learnable domain tokens projected by per-domain MLPs, producing domain-specific features $\mathbf{F}^{(k)}_{\rm SDM}$.
• Frequency-Adaptive Domain Fusion (FADF): At test time, lacking domain labels, the model computes frequency prototypes for each source domain and for the test image. Cosine similarity between these prototypes yields softmax weights $w_k$, used to fuse the domain-specific features: $$\mathbf{F}_{\rm fused} = \sum_k w_k \mathbf{F}_{\rm SDM}^{(k)}$$.
• Hierarchical Mask-Prompt Refiner (HMPR): Uses a two-stage decoder regime. Stage 1 yields a coarse mask $M_{256}$ at $256 \times 256$; this is re-encoded as a prompt, refined with self-attention, and supplied to Stage 2 to generate a finer mask $\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$0, eventually upsampled to $\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$1.

This architecture facilitates explicit disentanglement of frequency-domain information and dynamic domain adaptation during both training and inference (Wang et al., 9 Jan 2026).

2. Frequency Decomposition and Domain Modulation

Classically, a 2D discrete wavelet transform (DWT) produces low-frequency ($\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$2) and high-frequency ($\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$3) components using fixed filters $\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$4, $\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$5:

$\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$6

WaveRNet replaces fixed filters with parameterized convolutional branches:

$\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$7

$\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$8

These are fused as:

$\mathbf{F} \in \mathbb{R}^{C \times H'\times W'}$9

where $H' = H/16$0 is a learnable residual scalar. This preserves both illumination-robust low-frequency content and vessel-boundary high-frequency structures critical for segmentation.

SDM operates by associating each domain $H' = H/16$1 with a learnable token $H' = H/16$2, projected via a two-layer MLP:

$H' = H/16$3

This token is spatially broadcast and added to $H' = H/16$4 to produce $H' = H/16$5. During training, the token corresponding to the sample's domain is used, enabling domain awareness during representation learning (Wang et al., 9 Jan 2026).

3. Frequency-Adaptive Domain Fusion (FADF)

FADF addresses the problem of domain identification at test time, when no domain label is available. After training, per-domain prototypes are computed:

$H' = H/16$6

$H' = H/16$7

For a test image, features $H' = H/16$8, $H' = H/16$9 are extracted. Cosine similarity is measured for each source domain $W' = W/16$0:

$W' = W/16$1

Softmax weights (temperature $W' = W/16$2) yield the fusion coefficients:

$W' = W/16$3

$W' = W/16$4

This mechanism enables dynamic fusion of domain-adapted features based on frequency-domain similarity, constituting a frequency-driven soft domain selection (Wang et al., 9 Jan 2026).

4. Hierarchical Mask Generation and Refinement

The HMPR module mitigates detail loss observed in SAM's single-stage decoder. The two-stage process proceeds as:

• Stage 1: Decoder $W' = W/16$5 receives $W' = W/16$6 (spatial+prompt embeddings), producing mask $W' = W/16$7 at $W' = W/16$8 resolution.
• Prompt Feedback: $W' = W/16$9 is re-encoded as prompt $\mathbf{F}$0 via the SAM prompt encoder $\mathbf{F}$1, followed by self-attention to model long-range dependencies among vessel segments.
• Stage 2: Decoder $\mathbf{F}$2 takes $\mathbf{F}$3, outputs $\mathbf{F}$4 at $\mathbf{F}$5, which is bilinearly upsampled to $\mathbf{F}$6.

Iterative refinement with mask-based prompts and attention stages preserves fine vessel details, including capillaries and small branches commonly lost in direct upsampling workflows (Wang et al., 9 Jan 2026).

5. Optimization and Training Protocols

WaveRNet is trained via a composite objective:

$\mathbf{F}$7

with $\mathbf{F}$8, $\mathbf{F}$9. The loss terms are:

• Dice Loss: $\mathbf{F}^{(k)}_{\rm SDM}$0
• Focal Loss ($\mathbf{F}^{(k)}_{\rm SDM}$1): $\mathbf{F}^{(k)}_{\rm SDM}$2
• IoU-prediction MSE: $\mathbf{F}^{(k)}_{\rm SDM}$3, where $\mathbf{F}^{(k)}_{\rm SDM}$4 is the actual IoU, providing mask-level self-assessment.

Training follows the Leave-One-Domain-Out (LODO) protocol on DRIVE, STARE, CHASE_DB1, and RECOVERY-FA19 datasets. Metrics include Dice, IoU, and F1 (Wang et al., 9 Jan 2026).

6. Empirical Results and Ablation Analyses

WaveRNet exhibits high performance in both in-domain and cross-domain evaluation:

• In-domain: Matches or slightly surpasses state-of-the-art U-Net variants; e.g., Dice ≈ 80.46% (DRIVE).
• LODO cross-domain: U-Net variants degrade (Dice ≈ 25–40%), while SAM-based fine-tuning approaches plateau around 31–55%. WaveRNet achieves avg Dice ≈ 69.5% (best next, SAM-Med2D-FT, ≈ 60.1%). For individual domains under LODO: DRIVE (78.55%), STARE (81.06%), CHASE_DB1 (76.58%), RECOVERY-FA19 (41.75%).

Ablations validate the contributions of core modules:

• SDM provides the largest single gain under LODO (+6.79% Dice).
• Both low- and high-frequency branches are essential; omitting either degrades LODO performance substantially.
• HMPR is effective only when coupled with frequency-domain adaptation, supporting the necessity of joint frequency segmentation and hierarchical refinement.

7. Implementation Specifics

The model is implemented in PyTorch 2.1.0 with CUDA 12.8 and mixed-precision training on an NVIDIA RTX 4070 Ti GPU. Typical settings include batch size 2, 100 epochs, Adam optimizer (learning rate $\mathbf{F}^{(k)}_{\rm SDM}$5, exponential decay $\mathbf{F}^{(k)}_{\rm SDM}$6 per epoch), and initialization parameters ($\mathbf{F}^{(k)}_{\rm SDM}$7, $\mathbf{F}^{(k)}_{\rm SDM}$8, domain-token std=0.02). Image preprocessing replicates the official SAM pipeline (Wang et al., 9 Jan 2026).

The modular design and empirical validation of WaveRNet establish it as a reference method for robust domain-generalized segmentation in retinal imaging, providing a foundation for future frequency-based domain adaptation and fine-structure preservation strategies in medical image analysis.