TPOT: TOPOLOGY PRESERVING OPTIMAL TRANSPORT IN RETINAL FUNDUS IMAGE ENHANCEMENT

Abstract

Retinal fundus photography enhancement is important for diagnosing and monitoring retinal diseases. However, early approaches to retinal image enhancement, such as those based on Generative Adversarial Networks (GANs), often struggle to preserve the complex topological information of blood vessels, resulting in spurious or missing vessel structures. The persistence diagram, which captures topological features based on the persistence of topological structures under different filtrations, provides a promising way to represent the structure information. In this work, we propose a topology-preserving training paradigm that regularizes blood vessel structures by minimizing the differences of persistence diagrams. We call the resulting framework Topology Preserving Optimal Transport (TPOT). Experimental results on a large-scale dataset demonstrate the superiority of the proposed method compared to several state-of-the-art supervised and unsupervised techniques, both in terms of image quality and performance in the downstream blood vessel segmentation task. The code is available at https://github.com/Retinal-Research/TPOT.

1. INTRODUCTION

Retinal color fundus photography (CFP) plays a critical role in diagnosing ocular diseases [4, 5]. However, retinal fundus imaging equipment, particularly nonmydriatic cameras, faces challenges such as artifacts and blurring caused by uncontrollable factors (e.g., variability in ambient light conditions, capture errors, or operator mistakes). These issues often result in degraded retinal images. Some types of degradation, such as contrast mismatch, inconsistent illumination, and vessel occlusions, can impair the accurate visualization of lesions and blood vessels, essential for diagnosing retinal diseases such as diabetic retinopathy. Therefore, developing a robust framework to enhance low-quality images is vital for improving diagnostic accuracy [1, 2].

Since it is both difficult and costly to collect paired degraded and clean retinal images, research on retinal fundus image quality enhancement has shifted from supervised algorithms [6] to unsupervised algorithms, leading to remarkable improvements in image quality. These unsupervised approaches typically model retinal quality enhancement as an end-to-end image-to-image (I2I) translation task using Generative Adversarial Networks (GANs), where low-quality images $X$ are transformed into their high-quality counterparts $Y$. Specifically, CycleGAN [7] serves as a common strategy, as its cycle consistency regularization enables learning meaningful bidirectional mappings and improves generation alignment, albeit with additional computational cost and potential artifacts.

To address these limitations, Wang et al. [8] applied optimal transport (OT) theory to learn optimal mappings between two image domains, Zhu et al. [1, 2] introduced to use the Structural Similarity Index (SSIM) to regularize the generator and prevent excessive distortion of lesion structures. However, these methods often prove suboptimal, such as suffering from mode collapse, when handling images with complex structures or multimodal distributions. To mitigate this, Dong et al. [9] proposed the unpaired neural schrödinger bridge method to reserve the smooth and probabilistically consistent transformation between distributions. However, their iterative learning process can smooth high-frequency information, producing suboptimal structure preservation.

Topology preservation has become a key focus in segmentation tasks. Specifically, Hu et al. [10] extended beyond simple geometric structures by constructing a differentiable topology loss based on persistence diagrams and incorporating it into the segmentation training process. However, existing fundus image enhancement methods primarily focus on maintaining semantic consistency in the feature space [11] without considering introducing topology preservation and leveraging the natural advantage(e.g., vessel topology structures).

In this paper, we propose a novel topology-preserving retinal fundus image enhancement framework, named Topology Preserving Optimal Transport (TPOT), which leverages Optimal Transport (OT) theory to learn the optimal mappings between low-quality and high-quality images while incorporating topology regularization to preserve the complex blood vessel structures. Our main contributions are threefold: (i) To the best of our knowledge, this is the first work to introduce topology-preserving into OT-GAN-based retinal enhancement methods. (ii) We extend the application of topology loss beyond segmentation tasks, demonstrating its effectiveness in preserving topological structures during image generation. (iii) We conduct experiments on public retinal fundus datasets, showing our method outperforms various supervised and unsupervised approaches in denoising and downstream vessel segmentation.

2. TOPOLOGY PRESERVING OPTIMAL

Our proposed methods contain two main modules. Optimal Transport guided domain transformation and topology-preserving regularization for retinal blood vessels. The whole framework is shown in Fig. 1.

Fig. 1.

(A) illustrates the TPOT framework. The structures of $G_{\theta }$ and $D_{\beta }$ follows the design outlined in [1, 2], and $S_N$ follows the design presented in [3]. (B) represents the changes in segmentation masks of the images during training. The orange box highlights regions with complex topological structures. (C) provides an example of how our topology-preserving regularization operates. The yellow crosses and blue dots represent corresponding persistent features in $S_N(\mathbf{x})$ and $S_N\left(G_{\theta }(\mathbf{x})\right)$, respectively. The green dots indicate persistent noise features that must be removed during training. The model penalizes the differences between corresponding points to encourage topological consistency. Further details are discussed in Sec. 2.

Quality enhancement guided by optimal transport.

Let $\mathbf{X}$ and $\mathbf{Y}$ represent the domains of low-quality and high-quality images, with corresponding probability measures $\mu$ and $\nu$. The retinal fundus image enhancement task can be naturally framed as solving Monge’s optimal transport problem, which is formulated as follows:

$$\begin{gathered} f^{\mathrm{*}}=\underset{f}{\mathrm{i}\mathrm{n}\mathrm{f}}{\int }_{\mathbf{x}}C(\mathbf{x},f(\mathbf{x}))d\mu (\mathbf{x}) \\ \mathrm{s}\mathrm{u}\mathrm{b}\mathrm{j}\mathrm{e}\mathrm{c}\mathrm{t}\mathrm{t}\mathrm{o}{f}_{\#}\mu =\nu \end{gathered}$$ (1)

where $C(,\cdot ):\mathbf{X}\times \mathbf{Y}\to {\mathbb{R}}^{+}$ represents the cost function, $f$ denotes the candidate mapping functions, and the constraint ensures alignment between the two probability measures. Inspired by the works in [1, 2, 8], Eq. 1 can be modeled using an adversarial training strategy. Specifically, the generator $G_{\theta }$, parameterized by $\theta$, is used to model the optimal mapping $f^{\mathrm{*}}$. Identity regularization $C\left(\mathbf{y},G_{\theta }(\mathbf{y})\right)$ ensures that it doesn’t introduce unnecessary changes. The discriminator $D_{\beta }$ helps model the Wasserstein-1 distance between ${\mathbb{P}}_{\mathbf{Y}}$ and ${\mathbb{P}}_{G_k}(\mathbf{X})$, fulfilling the pushforward condition in Eq. 1. This leads to the following objective function:

$$\begin{array}{c}\underset{G_{\theta }}{\mathrm{m}\mathrm{a}\mathrm{x}}\underset{D_{\beta }}{\mathrm{m}\mathrm{i}\mathrm{n}}{\mathbb{E}}_{\mathbf{x}\sim {\mathbb{P}}_{\mathbf{x}}}\left[C\left(\mathbf{x},G_{\theta }(\mathbf{x})\right)\right]+{\mathbb{E}}_{\mathbf{y}\sim {\mathbb{P}}_{\mathbf{Y}}}\left[C\left(\mathbf{y},G_{\theta }(\mathbf{y})\right)\right]+\\ \lambda {\mathbf{W}}_1\left({\mathbb{P}}_{\mathbf{Y}},{\mathbb{P}}_{G_{\theta }(\mathbf{x})}\right)\end{array}$$ (2)

Here, Cost functions $C\left(\mathbf{x},G_{\theta }(\mathbf{x})\right)$ and $C\left(\mathbf{y},G_{\theta }(\mathbf{y})\right)$ are calculated based on the locally Quasi-Convex Multi-Scale Structural Similarity Index Measure [12, 13]. To enforce 1-Lipschitz continuity in the discriminator, a gradient penalty is applied.

Topology-preserving regularization.

Persistent homology provides a continuous representation of topological changes as the threshold $\alpha$ dynamically varies. Specifically, let $S_N$ denote the pre-trained segmentation network, $S_N(\mathbf{x})$ and $S_N\left(G_{\theta }(\mathbf{x})\right)$ represent the output likelihood map of low-quality input and synthetic high-quality counterparts, respectively. (e.g., $S_N(\mathbf{x}{)}^{0.5}=\left\{\mathbf{x}\mid S_N(\mathbf{x})>0.5\right\}$ represents the binary segmentation mask for $\mathbf{x}$) As $\alpha$ gradually decreases, an increasing number of pixels in $S_N{\left(G_{\theta }(\mathbf{x})\right)}^{\alpha }$ are classified as foreground (e.g., blood vessels), leading to corresponding changes in the topological structures. For instance, as $\alpha$ decreases, vessels that form loops may lose their structure as more neighboring pixels are considered foreground vessels. The persistence diagrams, which summarize the changes in topology, are denoted as $Dgm\left(S_N(\mathbf{x})\right)$ and $Dgm\left(S_N\left(G_{\theta }(\mathbf{x})\right)\right)$. These diagrams yield various persistent points, denoted as $P=\left(d_1,b_1\right)$ and $P^G=\left(d_2,b_2\right)$, where $d_1$ and $d_2$ correspond to the death thresholds and $b_1$ and $b_2$ correspond to the birth thresholds of the topological features. Following the ideal outlined in [10], the topology regularization, $C_{topo}\left(\mathbf{x},G_{\theta }(\mathbf{x})\right)$, is defined as:

$$\begin{gathered} C_{\mathit{topo}}\left(\mathbf{x},G_{\theta }\mid (\mathbf{x})\right)=\underset{\alpha }{\mathrm{m}\mathrm{i}\mathrm{n}}\sum _{P^G\in Dgm\left(S_N\left(G_{\theta }\right)\right)}{‖P^G-\alpha \left(P^G\right)‖}^2 \\ ={\sum }_{P^G\in Dgm\left(S_N\left(G_{\theta }\right)\right)}{\left(b_2-{\alpha }^{\mathrm{*}}\left(b_2\right)\right)}^2+{\left(d_2-{\alpha }^{\mathrm{*}}\left(d_2\right)\right)}^2 \end{gathered}$$ (3)

where $\left({\alpha }^{\mathrm{*}}\left(d_2\right),{\alpha }^{\mathrm{*}}\left(b_2\right)\right)\in P$ and ${\alpha }^{\mathrm{*}}$ represents the optimal correspondence of points between $P$ and $P^G$. This correspondence is calculated based on the persistence (e.g., $\mid d_1-b_2\mid$) of the critical points in $P$ and $P^G$. We compute the topology regularization based on image patches. The final topology regularization is expressed as:

$$C_{topo}\left(\mathbf{x},G_{\theta }(\mathbf{x})\right)=\sum _i{C}_{topo}^i\left({\mathbf{x}}_{\mathbf{i}},G_{\theta }\left({\mathbf{x}}_{\mathbf{i}}\right)\right)$$ (4)

where $i$ represents the index of the image patches. As illustrated in Fig. 1, by enforcing the alignment of persistent dots in $P^G$ with their corresponding dots in $P$ and by removing noisy dots that lack matches, we ensure that the enhancement process preserves the structural integrity of blood vessels. Consequently, as training progresses, the complex topological structure of the blood vessels (highlighted in the orange box) gradually becomes more consistent between $\mathbf{x}$ and $G_{\theta }(\mathbf{x})$. Additionally, since Eq. 3 is computed based on the mean squared error between critical points in $P$ and $P^G$, this regularization is convex and has a well-established gradient [10]. Thus, the final training objective is given by:

$$\begin{array}{l}\underset{G_{\theta }}{\mathrm{m}\mathrm{a}\mathrm{x}}\underset{D_{\beta }}{\mathrm{m}\mathrm{i}\mathrm{n}}{\lambda }_{\mathit{ssim}}{\mathbb{E}}_{\mathbf{x}\sim {\mathbb{P}}_{\mathbf{X}}}\left[C\left(\mathbf{x},G_{\theta }(\mathbf{x})\right)\right]+\\ {\lambda }_{idt}{\mathbb{E}}_{\mathbf{y}\sim {\mathbb{P}}_{\mathbf{Y}}}\left[C\left(\mathbf{y},G_{\theta }(\mathbf{y})\right)\right]+\\ {\lambda }_{\mathit{topo}}{\mathbb{E}}_{\mathbf{x}\sim {\mathbb{P}}_{\mathbf{X}}}\left[C_{\mathit{topo}}\left(\mathbf{x},G_{\theta }(\mathbf{x})\right)\right]{\mathbb{I}}_{t_i}+\\ {\mathbf{W}}_1\left({\mathbb{P}}_{\mathbf{Y}},{\mathbb{P}}_{G_{\theta }(\mathbf{X})}\right)\end{array}$$ (5)

where ${\mathbb{I}}_{t_i}$ is the indicator function, signifying that $G_{\theta }$ incorporates topology regularization after $t_i$ epochs.

3. EXPERIMENTS AND RESULTS

3.1. Experimental Details

Baselines

We compare our methods against the following baselines: Supervised algorithm: cofe-Net [6]; Unsupervised algorithms: CycleGAN [7], OTT-GAN [8], OTRE-GAN [1], WGAN [14], Context-aware OT [11] and CUNSB-RFIE [9]. For the quality enhancement task, we utilize Structural Similarity Index Measure (SSIM) and Peak Signal-to-Noise Ratio (PSNR) as evaluation metrics. For the downstream vessel segmentation task, we assess performance using Area under the ROC Curve (AUC), Area under the Precision-Recall Curve (PR), F1-score and Specificity (SP).

Dataset.

For the quality enhancement task, we used the algorithms outlined in [6] to generate degraded images that combined spot artifacts, illumination changes, and blurring. All models were trained on 5,000 EyeQ images [15] and tested on 6,217 EyeQ images. The training utilized both supervised and unsupervised modes, depending on the specific methods employed. The vessel segmentation task was conducted using the enhanced images, starting from scratch with the DRIVE [16] dataset, adhering to the official training and testing split. Before any operations, All images are cropped and resized into 256 × 256 pixels, except method [6].

TPOT settings.

The generator and discriminator used in Eq. 5 follows the design outlined in [1,2], and the pre-trained segmentation network employed from [3]. We keep the same parameter settings with [6, 2, 11, 9] for all baselines. For TPOT, we trained the model for 200 epochs using the RM-Sprop optimizer, with $t_i=100$. The learning rate in each phase started at 2 × 10⁻⁴ and decayed by a factor of 10 after every 50 epochs. Additional parameters were set as follows: ${\lambda }_{ssim}=30,{\lambda }_{idt}=15,{\lambda }_{topo}=1$ and the patch size in Eq. 4 was set to 65.

3.2. Experimental Results

Quality enhancement task.

The metric scores are shown in Tab. 1, where our methods achieve the best performance in both SSIM and PSNR, highlighting its superiority in preserving perceptual quality and pixel-wise alignment. For other models, as shown in Fig. 2(A), the Context-aware OT method shows noticeable semantic deficiencies at the image boundaries, resulting in the disappearance of blood vessels. In contrast, CUNSB-RFIE maintains better consistency in image space and effectively handles spot artifacts, though it slightly alters the background. Other GAN-based models (e.g., CycleGAN, OTRE-GAN, OTT-GAN) struggle to preserve peripheral blood vessels with complex topological structures (e.g., loops), still resulting in over-tampering of vessel structures. Our method outperforms these approaches by preserving overall image quality and vessel structures.

Table 1.

Performance comparison with baselines.

	Method	EyeQ		Downstream vessel segmentation
		SSIM ↑	PSNR↑	AUC ↑	PR ↑	F1-score ↑	SP↑
Supervised Methods	cofe-Net [6]	0.9408	24.907	0.9188	0.7698	0.6890	0.9801
Unsupervised Methods	Context-aware OT [11]	0.8497	21.088	0.8650	0.6547	0.5887	0.9765
	CycleGAN [7]	0.9313	25.076	0.9015	0.7277	0.6462	0.9801
	OTRE-GAN [1]	0.9392	24.812	0.9156	0.7678	0.6918	0.9797
	OTT-GAN [8]	0.9275	24.065	0.9034	0.7400	0.6608	0.9812
	WGAN [14]	0.9266	24.793	0.9081	0.7494	0.6768	0.9764
	CUNSB-RFIE [9]	0.9121	24.242	0.9163	0.7625	0.6872	0.9784
	Ours	0.9417	25.196	0.9191	0.7748	0.6926	0.9816

The best performance within each column is highlighted in bold.

Fig. 2.

(A) illustrates the result of quality enhancement task on the EyeQ dataset. (B) shows the enhanced DRIVE image along with the corresponding vessel segmentation. The orange box highlights a patch with a complex topological structure, where our method demonstrates greater consistency.

Vessel segmentation task.

As shown in Tab. 1, our method achieves the best performance across all four evaluation metrics. Specifically, Fig. 2(B) presents the vessel segmentation illustration on the DRIVE dataset. Due to vessel deficiencies at the image level, the Context-OT model produces noisy boundary predictions. Other GAN-based methods (e.g., OTRE-GAN and CycleGAN) improve overall image quality but struggle to preserve peripheral vessel structures. In the highlighted segmentation patch, CUNSB demonstrates a better ability to predict small vessel branches, while other methods tend to classify them as background. Our method preserves the complex topology structure, leading to continuous and consistent vessel morphology.

3.3. Ablation Study

Our method leverages topology-preserving regularization to maintain the blood vessel structure during enhancement. To evaluate the effect of this additional regularization on the final generation performance, we conducted ablation studies on the EyeQ dataset, with the results reported in Tab. 2. The absence of topology-preserving regularization, as represented by OTRE-GAN, resulted in relatively poorer performance. As we incrementally increased the weight ${\lambda }_{topo}$ in Eq. 5, we observed a clear performance improvement, even with a small weight. This demonstrates that topology-preserving regularization contributes significantly to the overall enhancement quality.

Table 2.

Ablation study over Topology-preserving regularization.

${\lambda }_{topo}$	SSIM	PSNR
0 (OTRE-GAN)	0.9392	24.812
0.1	0.9409	24.963
1	0.9417	25.196
10	0.9413	25.126

4. CONCLUSION

In this study, we propose a Topology Preserving Optimal Transport (TPOT) pipeline for enhancing retinal images. We incorporate topology regularization to preserve blood vessel structures, promoting vessel alignment in the enhancement process. We demonstrate the effectiveness of TPOT in denoising and maintaining the topological integrity of blood vessels on EyeQ and DRIVE datasets, respectively. While we utilized synthetic degraded images and faced challenges related to domain shifts in capturing segmentation outputs, our findings underscore the importance of preserving topological structures in medical image generation. Future research will further explore the potential of topology regularization in other medical image generation applications.