Multi-Scale Target-Aware Representation Learning for Fundus Image Enhancement

Published in Neural Networks (2025) — Online article link

🚀 Environment Setup

1. Clone the repository

   git clone https://github.com/RyanWu31/MTRL-FIE.git
   cd MTRL-FIE
   

2. Install dependencies

    conda create -n mtrl python=3.10
    conda activate mtrl
    pip install -r requirements.txt
   

🧠 Training and Evaluation

1. Train

    python train.py --dataroot ./images/BA_normal_dataset/ \
                    --name revise_ems_confidence \
                    --model maenet \
                    --direction AtoB \
                    --dataset_mode cataract_with_mask \
                    --norm instance \
                    --batch_size 16 \
                    --gpu_ids 3 \
                    --lr_policy linear \
                    --n_epochs 180 \
                    --n_epochs_decay 50 \
                    --input_nc 1 \
                    --display_port 8099

2. Test

   python test.py --dataroot ./images/BA_normal_dataset/ \
               --name revise_ems_confidence \
               --model maenet \
               --dataset_mode cataract_with_mask \
               --load_size 512 \
               --crop_size 512 \
               --input_nc 1
   
   

Architecture Overview

The proposed MTRL-FIE method processes medical images through wavelet embedding and multi-scale feature extraction:

Pipeline Description:

• Origin I: Input retinal image
• D(I): Initial enhancement through wavelet decomposition
• MFE: Multi-scale Feature Extraction for capturing fine details and large-scale anatomical structures
• SHD: Spatial-Hierarchical Decoder
• TFA: Task-aware Feature Aggregation
• Reconstructed I: Final enhanced output image

The method extracts both fine-grained features and large-scale anatomical information, then selectively aggregates them through the TFA module to produce the enhanced image.

📊 Quantitative Results

Table 1. Comparison of image enhancement quality across different methods on five datasets of varying scales.

Methods	BA (SSIM / PSNR)	DRIVE (SSIM / PSNR)	Kaggle (SSIM / PSNR)	Subset-EyeQ (SSIM / PSNR)	Refuge (SSIM / PSNR)
DCP (He et al., 2010)	0.864 ± 0.082 / 20.58 ± 8.08	0.885 ± 0.075 / 19.36 ± 5.80	0.867 ± 0.072 / 18.05 ± 4.15	0.846 ± 0.082 / 18.31 ± 4.72	0.815 ± 0.075 / 16.23 ± 2.94
Cofe-Net (Shen et al., 2020)	0.910 ± 0.057 / 20.54 ± 3.31	0.930 ± 0.029 / 22.10 ± 4.18	0.949 ± 0.030 / 25.61 ± 4.34	0.919 ± 0.042 / 24.87 ± 4.51	0.933 ± 0.031 / 25.82 ± 3.55
StillGAN (Ma et al., 2021)	0.913 ± 0.054 / 23.45 ± 5.90	0.944 ± 0.040 / 27.91 ± 6.53	0.947 ± 0.035 / 28.56 ± 9.31	0.931 ± 0.049 / 28.08 ± 8.63	0.925 ± 0.051 / 29.75 ± 8.52
ArcNet (Li et al., 2022)	0.903 ± 0.042 / 19.87 ± 3.55	0.943 ± 0.032 / 23.70 ± 2.56	0.921 ± 0.028 / 21.99 ± 2.76	0.900 ± 0.040 / 22.17 ± 2.98	0.893 ± 0.040 / 21.79 ± 2.48
RFormer (Deng et al., 2022)	0.906 ± 0.056 / 23.64 ± 7.40	0.925 ± 0.039 / 26.20 ± 5.04	0.926 ± 0.048 / 25.73 ± 6.70	0.899 ± 0.064 / 23.57 ± 6.68	0.912 ± 0.046 / 25.08 ± 5.28
GfeNet (Li et al., 2023b)	0.928 ± 0.029 / 22.65 ± 2.53	0.946 ± 0.026 / 25.59 ± 3.61	0.957 ± 0.020 / 27.10 ± 3.53	0.944 ± 0.034 / 27.10 ± 2.90	0.935 ± 0.040 / 27.22 ± 3.11
MTRL-FIE (Ours)	0.933 ± 0.024 / 22.97 ± 3.42	0.950 ± 0.025 / 26.05 ± 4.36	0.963 ± 0.019 / 30.49 ± 5.52	0.947 ± 0.029 / 28.13 ± 4.31	0.950 ± 0.025 / 29.05 ± 3.41

🙏 Acknowledgments

We are extremely grateful for the assistance and support provided by our collaborators.

https://www.linkedin.com/in/ghulam-mustafa-bme