Context-Aware Optimal Transport Learning for Retinal Fundus Image Enhancement

TL;DR

This paper proposes a context-informed optimal transport (OT) learning framework for tackling unpaired fundus image enhancement, and derives the proposed context-aware OT using the earth mover's distance and shows that the proposed context-OT has a solid theoretical guarantee.

Abstract

Retinal fundus photography offers a non-invasive way to diagnose and monitor a variety of retinal diseases, but is prone to inherent quality glitches arising from systemic imperfections or operator/patient-related factors. However, high-quality retinal images are crucial for carrying out accurate diagnoses and automated analyses. The fundus image enhancement is typically formulated as a distribution alignment problem, by finding a one-to-one mapping between a low-quality image and its high-quality counterpart. This paper proposes a context-informed optimal transport (OT) learning framework for tackling unpaired fundus image enhancement. In contrast to standard generative image enhancement methods, which struggle with handling contextual information (e.g., over-tampered local structures and unwanted artifacts), the proposed context-aware OT learning paradigm better preserves local structures and minimizes unwanted artifacts. Leveraging deep contextual features, we derive the proposed context-aware OT using the earth mover's distance and show that the proposed context-OT has a solid theoretical guarantee. Experimental results on a large-scale dataset demonstrate the superiority of the proposed method over several state-of-the-art supervised and unsupervised methods in terms of signal-to-noise ratio, structural similarity index, as well as two downstream tasks. The code is available at \url{https://github.com/Retinal-Research/Contextual-OT}.

Figures (7)

• Figure 1: Comparison between the traditional OT learning scheme (Left) and the proposed context-aware OT scheme (Right) for fundus image enhancement. Different from the traditional OT learning scheme that performs OT on image space, our contextual OT performs OT on the contextual feature space, which can help preserve contextual information between low-quality and high-quality images.
• Figure 2: The adversarial training scheme of the proposed contextual OT. The generator ($f_{\theta}$) is a U-Net with residual connection and channel attention as outlined in zhu2023otre9763342. The discriminator ($\Phi$) is also the one used in zhu2023otre9763342. We use a VGG-19 for encoding the contextual information onto feature space and compute the contextual transport cost based on these feature embeddings.
• Figure 3: Qualitative results of the proposed method on degraded images over the combinations of different noise (i.e., spot artifacts, illumination, and blurring). Our method achieves good enhancement performance even on severely degraded images (cols. 4, 5, and 6).
• Figure 4: Ablation study on different $\lambda$.
• Figure 5: Visual comparison of our method with baseline methods. Red box highlights Low Illumination introduced, and Green box highlights the light spot artifact noise.