WSRFNet: Wavelet-Based Scale-Specific Recurrent Feedback Network for Diabetic Retinopathy Lesion Segmentation

1 Introduction

Diabetic retinopathy (DR), which is the main ocular complication of diabetes mellitus, has been the major cause of visual impairment and blindness, especially in the working-age individuals today [Ting et al., 2016; Thomas et al., 2019]. It is estimated that about 103 million adults suffered from DR all over the world in 2020. The numbers are predicted to increase to 161 million in 2045 [Teo et al., 2021]. The common manifestations of DR lesions include hard exudates (EX), hemorrhages (HE), soft exudates (SE) and microaneurysms (MA). Ophthalmologists normally diagnose DR based on the presence of these lesions in colorful fundus images. Due to the small number of ophthalmologists, one ophthalmologist needs to screen a great deal of fundus images, which is a heavy workload. Therefore, automatic DR lesion segmentation (DRLS) is of vital importance to assist the diagnosis of DR and reduce the workload of ophthalmologists.

∗Corresponding author.

A big challenge in DRLS is that the diverse lesions are largely different in size. For clear explanation of this challenge, we count the lesion size in DDR training set [Li et al., 2019a] and draw its cumulative distribution function (CDF). As shown in Fig. 1(a), the largest 10% lesions are at least 74 times larger than the smallest 10% lesions. Visual lesion annotation in the fundus image (see Fig. 1(b)) also illustrates the large variation in lesion size. An effective method to tackle this challenge is to aggregate multiscale features extracted from a feature encoder, which is shown in Fig. 1(b). The performance of the method depends on two factors, i.e., the effectiveness of multi-scale feature fusion module and the quality of the input features of different scales. Most existing methods [Guo et al., 2019; Bo et al., 2022] focus on the former factor to develop sophisticated multi-scale feature fusion modules and alleviate the challenge to some extent. However, the latter factor has not yet been investigated in DRLS. Observing this, our work will explore the latter factor in DRLS for the first time by designing a refinement module to rectify the multi-scale features to improve the quality of them for better aggregated multi-scale feature representation (see Fig. 1(b)).

The refinement of the multi-scale features has recently been investigated in other dense prediction tasks such as salient object detection (SOD) and camouflaged object detection (COD) tasks. One kind of method [Wei et al., 2020a; Hu et al., 2023], which has been proved to be effective, is to leverage the aggregated multi-scale features from previous iteration as feedback to refine the features of different scales via designed iterative feedback methods. We deeply analyze the works of such kind of method and find two main issues: 1) Spatial information loss in refinement. Current methods use the irreversible downsampling operations such as strided convolution and bilinear interpolation to downsample the feedback features [Wei et al., 2020a] or the refined multi-scale features [Hu et al., 2023], which may lead to the loss of spatial information that is important to pixel-wise predictions. 2) Scale-agnostic feedback for features of different scales. Large-scale features from low-level layers of the encoder contain rich details but inaccurate semantic information. Smallscale features from high-level layers contain precise semantics but lack enough details. Therefore, large-scale features may need more semantic refinement, while small-scale features may require more detail refinement. However, previous approaches apply the same feedback (i.e., scale-agnostic feedback) for features of different scales without considering their characteristics, which is difficult to meet their different requirements for the refinement.

To address above-mentioned issues, in this paper, we propose a Wavelet-based Scale-specific Recurrent Feedback Network, named WSRFNet, which recursively utilizes feedback to refine the multi-scale features to acquire better aggregated multi-scale feature representation for accurate DRLS. Specifically, for features of each scale, we generate feedback from the fused multi-scale features in previous loop to refine them. To reduce computation, we employ downsampling operation to obtain feedback of the same scale as the features needed to be corrected instead of upsampling the multi-scale features to the scale of high-resolution feedback [Hu et al., 2023]. To avoid the loss of feedback information when downsampling, we devise a wavelet-based feedback pyramid module (WFPM), which produces initial feedback features for features of each scale based on Haar wavelet transform (HWT) that is a lossless information transformation method [Stankovic and Falkowski, 2003]. Considering diverse characteristics of different scale features, we design a scale-specific refinement module (SRM), which conducts refinement with scale-specific feedback obtained by merging initial feedback features with features requiring correction.

In summary, this work makes the following contributions:

• For the DRLS task, we design a Wavelet-based Scalespecific Recurrent Feedback Network (WSRFNet), which aims to obtain exquisite aggregated multi-scale feature representation by refining the multi-scale features in a recurrent feedback manner.
• We devise a wavelet-based feedback pyramid module (WFPM) to keep as much information as possible via Haar wavelet transform, when the feedback information flows into each level of the network.
• We propose a scale-specific refinement module (SRM), which generates scale-specific feedback to pointedly refine the multi-scale features.
• Comprehensive experiments on IDRiD [Porwal et al., 2020] and DDR [Li et al., 2019a] datasets demonstrate the superiority of our method over the state-of-the-art (SOTA) approaches.

2 RelatedWork

2.1 Diabetic Retinopathy Lesion Segmentation

Existing DRLS works can be categorized into two main classes: one-type and multi-type lesion segmentation. For the former one, one-type lesion segmentation result is predicted with one forward propagation [Dai et al., 2018; Guo et al., 2020; Liu et al., 2021a; Zhang et al., 2022]. For instance, SS-MAF [Zhang et al., 2022] designs an auxiliary super-resolution task, which brings in helpful detailed features for tiny EX lesion and boundaries detection.

For the latter one, a unified model is applied to segment multi-type of lesions simultaneously. For example, L-seg [Guo et al., 2019] designs a multi-scale feature weighted fusion method to handle the issue that small lesion regions could not response at high level of network. RTNet [Huang et al., 2022] proposes to capture the intra-class dependencies among multi-lesion and inter-class relations between multilesion and vessels. SAA [Bo et al., 2022] devises a scaleaware attention block to re-weight importance of multi-scale features dynamically to alleviate the inconsistent scale problem. M2MRF [Liu et al., 2023] proposes a novel feature reassembly operator to preserve subtle activations about small lesions as much as possible. [Guo et al., 2024] utilize a set of parallel fully convolutional neural networks (FCN) with different input scales that can better extract the features of lesions with different sizes and introduce the weight CNN that can better fuse the advantages of these FCN. Different from previous methods, we focus on the refinement of multi-scale features to address the large scale variation problem for better performance of multi-type lesion segmentation.

2.2 FeedbackMechanisms

Feedback mechanisms have been successfully applied in various computer vision tasks such as super-resolution [Han et al., 2018; Li et al., 2019b], SOD [Wei et al., 2020a] and COD [Hu et al., 2023]. Most of the feedback mechanisms are implemented using iterative strategy, allowing the network to utilize output information to rectify previous states. For example, [Li et al., 2019b] design an image super-resolution feedback network, which refines low-level representation with high-level information under a recurrent structure. F3Net [Wei et al., 2020a] devises a cascaded feedback decoder for SOD, which can feedback features of both high semantics and high resolutions to previous ones to refine them. HitNet [Hu et al., 2023] embeds a global feedback connection into the multi-scale framework via a feedback fusion block, which integrates the information from multi-scale outputs. However, previous feedback mechanisms for dense prediction tasks adopt irreversible downsampling operations and produce the same feedback for multi-scale features, which may limit the performance of the models as mentioned above (Section 1). To maintain as much information for correction as possible when downsampling, we apply HWT in our feedback generation module. To tailor the feedback to the features of each scale, we integrate the feedback with features requiring refinement to obtain scale-specific feedback.

3 Method

Fig. 2 illustrates the overall architecture of our proposed WSRFNet, which follows an encoder-decoder structure. The encoder is used to extract initial multi-scale features B = {Bi}5i=1 from the input images. Specifically, the encoder utilizes the backbone network {Ei}5i=1 to extract the multi-level features from the input images and passes the extracted features through two convolutional layers to adjust the number of their channels to C = 64 for subsequent processing [Wei et al., 2020b]. For ease of statement, i ∈ {1, 2, 3, 4, 5} indicates the i-th level from bottom to top. The output features Bi from different levels of the encoder correspond to different scales, i.e., 1 2i resolution H ×W of the input images.

The decoder, which consists of three components: WFPM, SRM and multi-scale feature aggregator (MFA), is designed to gain a high-quality fused multi-scale feature representation by refining features of different scales in a recurrent way. The details of recurrent refinement, WFPM, SRM, MFA and loss function are introduced in Section 3.1, Section 3.2, Section 3.3, Section 3.4 and Section 3.5, respectively.

3.1 Recurrent Refinement of WSRFNet

Most existing remedies to large size variation problem in DRLS is to develop delicate multi-scale feature fusion modules to strengthen integrated multi-scale feature representation. Differently, we achieve this goal from a novel perspective, i.e., the refinement of multi-scale features. We perform refinement using the recursive feedback mechanism, which is inspired by the effectiveness of the iterative feedback method used for the correction of multi-scale features in other dense prediction tasks [Wei et al., 2020a; Hu et al., 2023].

To be specific, the loop number n=1 assumes the first loop and no feedback features from previous loop. In the first loop, the initial multi-scale features are fed into MFA to fuse the multi-scale features. For loop number (n>1), WFPM starts by producing initial feedback pyramid F n = {F n i }5i=1 from the fused multi-scale features Sn-1 of previous loop for initial multi-scale features. Then, SRM utilizes F n to generate scale-specific feedback to refine the initial multiscale features. After that, the refined multi-scale features Bn = {Bn i }5i=1 are passed through the MFA same as the first loop. In each loop, the decoder shares weights. The above recurrent refinement can be formulated as

3.2 Wavelet-Based Feedback PyramidModule

Previous feedback methods employ irreversible downsampling operations such as bilinear interpolation to obtain feedback for initial multi-scale features, which consequentially leads to information loss. The lost information may contain spatial details of the objects which are important to high performance DRLS. To avoid information dropping, we design a WFPM to produce initial feedback pyramid from aggregated multi-scale features of previous loop with the assist of HWT (see Fig. 2(a)). HWT is a reversible downsampling operation because of its biorthogonal property. Therefore, incorporated with HWT, WFPM enables less information loss when generating initial multi-scale feedback.

Concretely, given that the resolution of the lowest-level initial features B1 is the same as Sn-1 , a 3×3 convolutional layer is used to enhance Sn-1 to produce the initial feedback features F n 1 ∈ R H2 × W2 ×C of the lowest level (i=1). For level number (i>1), initial low-level feedback features F n i-1 are passed through a 1×1 convolutional layer Conv1×1 to adjust their channel number to C for reduction of computation, if their channel number is 4C.

Then, HWT is applied to downsample the channelcompressed feedback features, nondestructively. Formally, take 2D feature map Z ∈ R H 2i-1 × W 2i-1 ×C (i>1) as the input of HWT for the illustration of the implementation of HWT. HWT employs four filters WLL, WLH, WHL, WHH to decompose Z into four Haar wavelet subbands, which are corresponding to ZLL ∈ R H 2i × W2i ×C , ZLH ∈ R H 2i × W2i ×C , ZHL ∈ R H 2i × W2i ×C , ZHH ∈ R H 2i × W2i ×C , respectively. ZLL contains all low-frequency information of Z. ZLH, ZHL and ZHH include all high-frequency information of Z. Next, to retain all the information of the input feature map, we utilize a cross-channel concatenation operation to merge all the wavelet subbands to obtain the output of the HWT, which is [ZLL, ZLH, ZHL, ZHH] ∈ R H 2i × W2i ×4C .

At last, a 3×3 convolutional layer Conv3×3 is used to enlarge receptive field to adapt the increasement of the level to get the initial feedback features F n i ∈ R H 2i × W2i ×4C for the i-th level. The WFPM can be formulated as

3.3 Scale-Specific Refinement Module

Large-scale initial features lack accurate semantics, which is caused by the small receptive field of the low-level layer. Small-scale initial features are coarse in details, due to multiple downsampling in the backbone network. The same feedback used by current works [Wei et al., 2020a; Hu et al., 2023] is hard to satisfy diverse needs of different scale features for the correction. Considering the unique characteristics of initial features of each scale, SRM aims to pointedly perform refinement via scale-specific feedback features. Fig. 2(b) shows the structure of our SRM.

To be specific, SRM contains five blocks, which are denoted by {SRMi}5i=1 . Each block is used for the refinement of features of each scale. For the implementation of SRMi, firstly, we combine initial feedback features F n i with i-th level initial features Bi via a cross-channel concatenation operation Concat. Then, the combined features are fed into a 1×1 convolutional layer to acquire scale-specific feedback features Xn i ∈ R H 2i × W2i ×C , which is calculated by

The 1×1 convolutional layer has three main purposes: 1) to compress the channel number of the combined features from 5C to C to align with the i-th level initial features Bi for subsequent refinement, 2) to adequately merge cross-channel information of the combined features so that characteristics of the i-th level initial features can be injected into the initial feedback to generate more targeted feedback for i-th level initial features, and 3) to filter redundant information to obtain more representative feedback.

After that, we rectify the initial features with scale-specific feedback features via an element-wise addition operation (⊕)

followed by a 3×3 convolutional layer for enhancement to get the refined features Bn i , which is denoted by

3.4 Multi-Scale Feature Aggregator

MFA aims to obtain multi-scale information by fusing the features of different scales via an FPN-like structure [Lin et al., 2017], which is shown in Fig. 2(a). In practice, MFA firstly utilizes a bilinear interpolation operation to upsample the high-level features Dn i by a factor of 2. Then, the output features are integrated with the low-level features Bn i-1 via an element-wise addition operation. After that, a 3×3 convolutional layer is applied to enhance the merged features. Taking Bn as input, MFA can be formulated as

3.5 Loss Function

To guarantee better fusion of the multi-scale features and get the satisfactory feedback, following [Hu et al., 2023], we impose constraint from the ground truth G on the segmentation maps predicted by the aggregated multi-scale features in each loop. In practice, {Sn }Nn=1 are fed into a 3×3 convolutional layer and an upsampling layer to produce the segmentation map {P n }Nn=1 , where N is the total loop number. The final loss function L is calculated by

where Ld is Dice loss [Milletari et al., 2016] that is commonly used in DRLS, nN represents the weight of each loop constraint, which is set similar to [Hu et al., 2023].

As the fused multi-scale features are supervised in each loop, they may contain critical information for better prediction. Noting this, we combine the aggregated multi-scale features of all loops for prediction via a weighted average operation in testing, which is expressed as

where n is the weight of the fused multi-scale features of each loop, which is set according to the weight of each loop constraint in training (Eq. 6), Sf are the final features for prediction. In testing, we pass Sf through the trained 3×3 convolutional layer and an upsampling layer to get the ultimate segmentation map Pf .

4 Experiments

4.1 Datasets and EvaluationMetrics

The performance of our method is evaluated on two popular benchmark datasets, i.e., IDRiD [Porwal et al., 2020] and DDR [Li et al., 2019a]. Both datasets contain four types of annotation, including EX, HE, SE and MA.

IDRiD Dataset

This dataset is available for DRLS in ISBI-2018 grand challenge [Porwal et al., 2018; Porwal et al., 2020], which has 81 colour fundus images with 54/27 images for training and testing. Each image contains 4288×2848 pixels.

DDR Dataset

This dataset is provided by Ocular Disease Intelligent Recognition (ODIR-2019) [Li et al., 2019a], which contains 757 colour fundus images with 383/149/225 images for training, validation and testing. The sizes of images range from 1088×1920 to 3456×5184 pixels.

Evaluation Metrics

To quantitatively measure the proposed model, we follow the suggestion proposed by [Liu et al., 2023], and adopt standard metrics including class-wise Area Under PrecisionRecall curve (AUPR), mean class-wise AUPR (mAUPR), class-wise Dice coefficient (Dice), mean class-wise Dice coefficient (mDice), class-wise intersection-over-union (IoU) and mean class-wise IoU (mIoU).

4.2 Implementation Details

EfficientNet-B1 [Tan and Le, 2019], pre-trained on ImageNet [Russakovsky et al., 2015], is used as the backbone network. The number of feature channels for the five levels of EfficientNet-B1 are 16, 24, 40, 112 and 1280. Rotation (90◦ , 180◦ and 270◦ ), flipping (horizontal and vertical) and random cropping are applied to augment the images. Following [Liu et al., 2023; Guo et al., 2019; Bo et al., 2022], each image is resized to 1440×960 pixels in IDRiD and 1024×1024 pixels in DDR during training. Following [Liu et al., 2023], we adopt poly strategy with the initial learning rate of 0.01 and power of 0.9. Our network is trained end-to-end using stochastic gradient descent (SGD) optimizer. Momentum and weight decay are set to 0.9 and 0.0005, respectively. We train the model for 80k iterations on IDRiD and 100k iterations on DDR with a mini-batch size of 4. All the experiments of our method are performed using one NVIDIA GeForce RTX 3090 GPU with 24GB of memory.

4.3 Comparison with State-of-the-art Methods

We quantitatively compare the proposed WSRFNet with eight SOTA methods, including L-seg [Guo et al., 2019], RTNet [Huang et al., 2022], SAA [Bo et al., 2022], Guo et al. [Guo et al., 2024], M2MRF-C [Liu et al., 2023], Deeplabv3+ [Chen et al., 2018], HRNet48 [Wang et al., 2020] and Swinbase [Liu et al., 2021b]. The last three approaches, which are for general semantic segmentation (GSS), are implemented using the same training setting as ours for a fair comparison. HRNet48 utilizes FCN [Long et al., 2015] as the decoder. Deeplabv3+ takes ResNet-50 [He et al., 2016] as the backbone network. Swin-base uses UPerNet [Xiao et al., 2018] as the decoder. The results of rest methods, which are specially designed for DRLS, are directly taken from original papers. A few details of some DRLS methods should be noted that: 1) RTNet utilizes two auxiliary datasets, which are DRIVE [Staal et al., 2004] and STARE [Hoover et al., 2000] with pixel-level vessel labels. 2) The backbone network of SAA is pretrained on the DDR training set.

Results on IDRiD Dataset

We show the quantitative comparison on IDRiD dataset in Table 1. Compared with the both GSS and DRLS methods, our WSRFNet achieves the best mAUPR, mDice and mIoU values. In the individual evaluation, we gain the best result on SE segmentation under all metrics.

Results on DDR Test Set

The comparative results of each method on DDR test set are listed in Table 2. This dataset is more challenging, for the reason that its test set contains more images. Although some previous methods can perform well on IDRiD dataset, most of them achieve unsatisfactory results on this dataset. It can be seen from the quantitative results that our WSRFNet also achieves the best performance among all other methods under mAUPR, mDice and mIoU metrics. Specifically, compared with the second-best method (i.e., [Guo et al., 2024]) on IDRiD dataset, we outperform their model on DDR test set by a large margin, i.e., 14.77% improvement on mAUPR and 12.86% improvement on mDice. It is worth noting that we obtain the best performance on HE and SE segmentation under all metrics in the individual evaluation.

4.4 Ablation Study

Before exploring the influence of each module, a hyperparameter (i.e., N) needs to be determined. N represents the total loop number of the decoder. Considering that the IDRiD dataset does not contain a validation set, following [Liu et al., 2023], hyperparameter tuning is performed on the validation set of DDR. We gradually increase N from 1 to 5 and measure the performance of corresponding models in Table 3. With the increase of N, the performance of our model gradually rises first, and then falls at N=5. One possible reason is that our model is locally optimal at N=4. As N increases, we may find a value of N for the globally optimal model. However, the computational cost will rise alongside the increase in N. To balance the computational cost and performance, we choose N=4 as default in our WSRFNet. Our baseline model (1st row) is the model with N set to 1. All the following investigations are performed on the DDR test set using three evaluation metrics, i.e., mAUPR, mDice and mIoU.

Exploration of the Recurrent Refinement

We probe into the influence of the recurrent refinement on the aggregated multi-scale feature representation in this section. Concretely, we conduct the exploration by comparing the performance of prediction P n produced by the fused multi-scale features of each loop. The results are listed in the first 4 rows of Table 4. The performance under mDice and mIoU metrics gradually rises with the increase of the loop number n, which illustrates that aggregated representation related to precise region coverage iteratively corrects itself under the recurrent refinement. The best mIoU and mDice values are in the 4th loop while the best mAUPR value is in the 3rd loop, which demonstrates that the fused multi-scale features may contain diverse crucial information for high performance DRLS in the different loops. Therefore, in testing, the weighted combination of the fused multi-scale features in all loops (Eq. 7) can enhance the corrected feature representation in the last loop, which is quantitatively proved by comparing the performance of P 4 (4th row) and Pf (5th row).

Exploration of the WFPM

We study the impact of the WFPM by directly replacing the WFPM in our full model with other downsampling operations used in other feedback networks, i.e., bilinear interpolation used by [Wei et al., 2020a] and stride convolution used by [Hu et al., 2023]. Specifically, the kernel size and stride of stride convolution are set to 3×3 and 2, respectively. It can be seen from the Table 5 that our WFPM achieves the best performance, which indicates the effectiveness of the WFPM.

Fig. 3 shows visual comparison of other methods and our WFPM. It can be observed that our approach can produce better results for small lesions, which is benefited from the ability of WFPM to retain as much information as possible.

Exploration of the SRM

We investigate the impact of the SRM by replacing the SRM in our full model with other refinement modules employed in other feedback methods, i.e., element-wise addition used by [Wei et al., 2020a] and cross-channel concatenation used by [Hu et al., 2023]. In practice, the refinement module based on element-wise addition is implemented by replacing Xn i in Eq. 4 with F n i . We conduct the cross-channel concatenation approach to refine the initial multi-scale features by removing Bi in Eq. 4. It can be observed from the Table 6 that our approach outperforms other methods, which demonstrates the effectiveness of the proposed SRM.

Visual results of other methods and ours are displayed in Fig. 4. We observe that our model can generate more accurate segmentation results for lesions of different scales. We provide an unsuccessful case marked with dotted boxes, which is shown in Fig. 4(d). This region is also incorrectly identified by other approaches. A possible reason is that the region is too similar to MA in appearance to properly segment it.

Extension to Other Backbones

To further demonstrate the effectiveness of the proposed refinement method, we combine it with other widely used backbones including VGG-16 [Simonyan and Zisserman, 2014], ResNet-50 [He et al., 2016] and EfficientNet-B0 [Tan and Le,

2019]. We take backbones with MFA as the baseline models. As shown in Table 7, the baseline models are significantly improved by the proposed method with acceptable increasement of number of parameters, which indicates that the proposed methods are effective to other popular backbones. It is worth noting that compared with other models using the same backbone network including L-seg (VGG-16), Deeplabv3+ (ResNet-50), and SAA (EfficientNet-B0) in Table 2, our approaches outperform them by large margins.

5 Conclusion

In this work, our goal is to deal with the challenge of large size variation in DRLS via improving the quality of multiscale features to augment the aggregated multi-scale feature representation. To achieve this, we propose a WSRFNet, which adopts a recurrent feedback manner to rectify multiscale features. In particular, to alleviate the issue of information loss in previous feedback models, we develop a WFPM to produce initial feedback for features of different scales via a lossless dowmsampling operation, i.e., HWT. To cater to the diverse requirements of different scale features for correction as much as possible, we devise the SRM. This module conducts refinement using scale-specific feedback, which is obtained by combing initial feedback with features requiring correction. Comprehensive experiments on two public benchmark datasets demonstrate the superiority of our WSRFNet over the state-of-the-art approaches.

Acknowledgments

This work was supported in part by the Natural Science Foundation of China under Grant 62073105, U23A20389 and 62306095, by the Natural Science Foundation of Heilongjiang Province of China under Grant ZD2022F002, and by the Heilongjiang Touyan Innovation Team Program.