Source-linked AI summary
DoFE: Domain-oriented Feature Embedding for Generalizable Fundus Image Segmentation on Unseen Datasets
Shujun Wang, Lequan Yu, Kang Li, Xin Yang, Chi-Wing Fu, Pheng-Ann Heng
TL;DR
Fundus segmentation networks can lose accuracy on unseen datasets because clinical images differ across domains. DoFE learns multi-source domain knowledge, embeds similarity-weighted domain features into semantic features, and reports stronger unseen-domain segmentation across optic cup/disc and vessel tasks.
Problem
Clinical fundus images exhibit domain shifts in appearance and image quality, limiting CNN generalization from source datasets to unseen target datasets.
Method
DoFE stores multi-source domain prototypes in a Domain Knowledge Pool and dynamically embeds similarity-induced domain features into image semantic features.
Results
DoFE generates satisfying unseen-dataset results across optic cup/disc and vessel segmentation and outperforms prior domain generalization and network regularization methods.
Takeaways & Limitations
Multi-source domain knowledge can be used to enrich semantic features for generalizable fundus image segmentation on unseen datasets.
Takeaways & Limitations
DoFE requires multiple annotated source-domain datasets, whose professional ophthalmologist annotations are expensive and difficult to acquire.
Abstract
from arXiv · showhide
Deep convolutional neural networks have significantly boosted the performance of fundus image segmentation when test datasets have the same distribution as the training datasets. However, in clinical practice, medical images often exhibit variations in appearance for various reasons, e.g., different scanner vendors and image quality. These distribution discrepancies could lead the deep networks to over-fit on the training datasets and lack generalization ability on the unseen test datasets. To alleviate this issue, we present a novel Domain-oriented Feature Embedding (DoFE) framework to improve the generalization ability of CNNs on unseen target domains by exploring the knowledge from multiple source domains. Our DoFE framework dynamically enriches the image features with additional domain prior knowledge learned from multi-source domains to make the semantic features more discriminative. Specifically, we introduce a Domain Knowledge Pool to learn and memorize the prior information extracted from multi-source domains. Then the original image features are augmented with domain-oriented aggregated features, which are induced from the knowledge pool based on the similarity between the input image and multi-source domain images. We further design a novel domain code prediction branch to infer this similarity and employ an attention-guided mechanism to dynamically combine the aggregated features with the semantic features. We comprehensively evaluate our DoFE framework on two fundus image segmentation tasks, including the optic cup and disc segmentation and vessel segmentation. Our DoFE framework generates satisfying segmentation results on unseen datasets and surpasses other domain generalization and network regularization methods.
I. INTRODUCTION
The paper targets poor generalization of fundus segmentation networks to unseen domains caused by clinical distribution shifts. DoFE uses multi-source domain knowledge to enrich features and improve segmentation on unseen datasets.
- Clinical fundus images vary across institutions, scanners, patient populations, disease severity, field of view, appearance, and image quality.
- Training on multiple source domains is intended to produce segmentation networks that generalize to unseen target domains without using target-domain images during training.
- DoFE dynamically augments semantic image features with domain prior knowledge learned from multiple source domains.
- The Domain Knowledge Pool stores source-domain prior information and supplies aggregated features according to similarity between the input image and source-domain images.
- A domain code prediction branch automatically learns domain relationships, while the framework dynamically embeds domain-oriented features with semantic features.
- Experiments cover optic cup/disc and vessel segmentation, with results reported as outperforming previous domain generalization and network regularization methods.
II. RELATED WORK
Related work spans retinal fundus segmentation, network regularization, domain adaptation, and domain generalization. These approaches seek stronger segmentation or improved robustness across domains through feature, input, output, or representation-level strategies.
- The paper reviews these techniques as context for improving generalization in retinal fundus segmentation.
- Retinal fundus segmentation has used handcrafted gradient, shape, texture, and boundary features alongside high-level CNN features.
- Network regularization methods include Mixup, Manifold Mixup, and CutMix, which train with interpolated images, hidden representations, or pasted image patches.
- Medical-image domain adaptation aligns source and target domains in input, feature, or output spaces.
- Domain generalization learns a representation from source domains for unseen targets without training on target images or labels.
III. METHODOLOGY
DoFE extends a multi-source segmentation network with a domain knowledge pool that stores source-domain prototypes and supplies similarity-weighted features. These features dynamically enrich semantic representations for unseen-domain segmentation.
- DoFE uses annotated source datasets and an unseen target dataset, training without any information from the target dataset.
- The framework addresses limited unseen-domain feature expressiveness by integrating a memory module into a vanilla segmentation network.
- 1) Domain knowledge pool:: Each knowledge-pool item represents a source dataset’s domain prior knowledge, initialized with its discriminative prototype.
- 1) Domain knowledge pool:: The prototype is formed by averaging high-level semantic features across spatial locations and images in each source dataset.
- 1) Domain knowledge pool:: The knowledge pool is updated during training with momentum to obtain more discriminative domain representations.
- The basic segmentation framework is DeepLabV3+ with a MobileNetV2 backbone, combining low-level and high-level features for fine-grained segmentation.
- For each input image, DoFE composes a domain-oriented aggregated feature from the pool and embeds it into the semantic feature to increase discrimination.
2) Domain similarity learning:
DoFE learns source-domain similarity with a domain code prediction branch and uses the predicted code to form the aggregated domain feature. The resulting representation supports domain-oriented feature embedding.
- The domain code prediction branch automatically learns the relationship between input images and source domains alongside the segmentation network.
- The branch predicts a normalized domain code from high-level global features using global average pooling, normalization, ReLU, convolution, and Softmax.
- The domain-oriented aggregated feature is formulated as a weighted sum of domain-knowledge-pool items according to the domain code.
3) Domain-oriented aggregated feature:
The framework aggregates domain knowledge into a feature map and combines it with semantic features through attention-guided selection. This dynamic embedding augments training features and recalibrates target-domain features during testing.
- The domain-oriented aggregated feature is formed from domain knowledge and tiled to match the semantic feature shape.The resulting feature hagg has spatial dimensions H × W and channel dimension C before tiling.
- An attention-guided mechanism dynamically augments semantic feature hs with the tiled aggregated feature ˆhagg.A convolutional layer followed by tanh generates a self-attention map that acts as a selective mask.
- The final domain-oriented feature h is passed to later segmentation components to generate fine-grained masks.Element-wise multiplication and addition combine the feature components before segmentation.
- During training, the embedding acts as automatic feature augmentation, while testing incorporates multi-source knowledge to recalibrate target-domain features.The two phases use the same scheme for regularization during training and feature discrimination during testing.
B. Learning Strategy
The domain code prediction branch learns source-domain identity as a classification-related signal while smoothing one-hot codes to reduce overfitting. Its training uses mean squared error, alongside supervised segmentation learning.
- The domain code prediction branch treats source-domain identification as a classification problem using domain attributes as supervision.Images from the same domain share the same domain code, making direct regression prone to overfitting.
- Domain code smooth: Hard one-hot domain codes are randomly smoothed by perturbing the target value into [0.8, 1.0] and assigning non-negative values to other entries.The smoothed values satisfy the stated constraint while preserving the source-domain target.
- Domain code smooth: The smoothed domain-code prediction is optimized as a regression problem using Mean Square Error.The loss is denoted Ldc, with predicted codes and smoothed ground truths defined for source-domain images.
- Segmentation loss: OD and OC segmentation use binary Cross-Entropy loss under the multi-label setting.This supervised segmentation loss complements the domain-code objective.
2) Total objective function:
The framework combines supervised segmentation and domain-code prediction losses into one training objective. Training initializes the knowledge pool from semantic features and uses multi-source batches to learn domain similarity and update the pool.
- Total objective function: The whole framework is optimized by combining domain code prediction loss Ldc with supervised segmentation loss Ls.The combined objective is written as L = Ls + αLdc.
- Total objective function: The balancing weight α is empirically set to 0.1 in the experiments.
- Training procedure: A vanilla segmentation network is first trained, then its domain-specific semantic features initialize the domain knowledge pool.The initialized pool is used when training the complete framework.
- Training procedure: Randomly sampled images from multiple source domains form diverse batches treated as pseudo-test images for learning similarity and updating the knowledge pool.Embedded target-domain features are then intended to enhance segmentation results.
IV. EXPERIMENTS
The experiments evaluate unseen-domain fundus segmentation across OC/OD and vessel tasks using datasets acquired with different scanners and institutions. They report task-specific metrics, repeated-run variability, and quantitative OC/OD comparisons against generalization methods.
- Datasets and experiment setting: The evaluation covers OC/OD and vessel segmentation across four public datasets captured with different scanners and institutions.Each dataset is treated as a domain, and OC/OD feature visualization shows separation among datasets.
- Evaluation metrics: The OC/OD task is evaluated with DSC, 95% Hausdorff Distance, and Average Surface Distance.DSC is defined from pixel-level true positives, false positives, and false negatives.
- Evaluation metrics: The vessel task uses accuracy, specificity, sensitivity, and ROC AUC as evaluation metrics.
- Evaluation protocol: Each experiment is run three times, with mean performance and standard deviation reported to account for training randomness.
- Quantitative comparison: Table II compares network and domain generalization methods on OC/OD segmentation using three-run means and standard deviations.Domain 1 results use models trained on the other three domains, and top results are bolded.
C. Implementation Details
The experiments compare DoFE with a unified DeepLabV3+ baseline, network regularization methods, and domain generalization methods for unseen-domain OC/OD segmentation. DoFE improves average DSC and segmentation distances, while qualitative results show more accurate boundaries in difficult images.
- Baseline comparison: DoFE surpasses a unified vanilla DeepLabV3+ baseline by 2.90% average DSC on OC/OD segmentation.
- Comparison with network regularization methods: CutMix reaches 86.62% average DSC, 1.08% above the baseline, while M-mixup reaches 86.41% and mixup performs worse than baseline.The authors attribute the stronger regularization results to better preservation of structure information than image-level interpolation.
- Comparison with domain generalization methods: DST raises average DSC from 85.54% to 86.83%, but DoFE still achieves a further 1.61% average DSC improvement without stacked augmentation.
- Qualitative results: Figure 4 compares original images and enlarged ROIs across four domains using predicted optic-disc, optic-cup, and ground-truth contours.Green and blue contours mark optic discs and cups, respectively; red contours mark ground truth.
- Comparison with domain generalization methods: DoFE outperforms JiGen by 1.47% average DSC and produces the smallest average ASD and HD errors, except for OD HD, which is comparable with DST.
- Qualitative results: In low-contrast and abnormal samples, DoFE produces more accurate optic-disc and optic-cup boundaries than the other methods.
4) Qualitative results:
Vessel-segmentation experiments compare DoFE with a vanilla DeepLabV3+ baseline across four test domains, while the ablation study examines individual framework components. The evaluation also notes a boundary on comparisons with prior supervised methods.
- E. Experiment Results on Vessel Segmentation: Across four vessel-segmentation test domains, DoFE achieves more accurate results than the vanilla DeepLabV3+ baseline on ACC, SP, and AUC.For each test case, the other three datasets serve as training data.
- E. Experiment Results on Vessel Segmentation: Table IV reports quantitative vessel-segmentation comparisons across four test domains using percentage metrics.
- E. Experiment Results on Vessel Segmentation: Absolute vessel-segmentation performance should not be directly compared with previous supervised methods because the backbones and experimental settings differ.
- 1) Ablation study: The ablation study removes or freezes the knowledge pool, domain-code prediction, self-attention guidance, and related feature-embedding components.The variants are labeled w/o KP, w/o DC, w/o SM, and w/o Tr.
1) Ablation study:
Additional experiments test domain-oriented feature embedding against random feature augmentation, domain adaptation, and component variants. The results support learned domain knowledge, report statistically clear improvements, and indicate low added inference cost.
- Ablation study: DoFE outperforms random Gaussian feature augmentation on four datasets, indicating a positive impact from learned domain knowledge.Both methods augment semantic features during training and testing, but DoFE uses a learned domain knowledge pool.
- Domain adaptation comparison: The domain-adaptation experiment uses target-domain images during training and compares DoFE with the state-of-the-art BEAL method.The multi-source training datasets are merged and treated as one source domain.
- Statistical analysis: Paired t-tests use a 0.05 significance level and 95% confidence level, with clear improvements for DoFE over other methods and its designed components.
- Ablation study: Table V defines ablation variants that remove knowledge-pool embedding, domain-code prediction, self-attention guidance, or knowledge-pool training.
- Computation cost: DoFE adds only 0.03 seconds of average inference time per image, increasing vanilla DeepLabV3+ time from 0.70s to about 0.73s.The framework adds two fully connected layers and one convolution layer.
5) Computation cost:
DoFE uses multi-source domain knowledge to improve fundus segmentation generalization, but its current formulation remains constrained by data and task-setting requirements.
- Limitations: DoFE requires multiple annotated source-domain datasets, whose collection and professional annotation are expensive.The authors identify unlabeled source datasets and semi-supervised learning as a potential solution.
- Limitations: The current method considers general discrepancy between datasets but not inter-domain or intra-domain differences involving disease type and severity.The authors propose addressing these discrepancies in future work.
- Limitations: The current problem setting requires source and target domains to share the same segmentation task.The authors plan to study memory-based general transfer learning for different tasks.
- Method: DoFE incorporates a domain knowledge pool to learn and memorize prior information from multiple source datasets.Domain features are aggregated according to similarity between testing and source domains, then embedded with semantic features using self-attention.
- Evaluation: The framework was evaluated on two retinal fundus image segmentation tasks and demonstrated improvements in effectiveness.The tasks cover optic cup/disc segmentation and retinal vessel segmentation.