Coarse Retinal Lesion Annotations Refinement via Prototypical Learning

The input to the network is the concatenation of the image patch $I\in R^{H\times W\times 3}$ and its corresponding coarse lesion annotation $M\in R^{H\times W\times 1}$. We use a modified U-Net backbone to extract its feature map $F\in R^{H^{\prime }\times W^{\prime }\times C}$. Following [14], we remove the last two upsampling blocks in the U-Net to speed up the calculation. As a result, the resolution of the feature map is 1/4 of the original input. We concatenate the feature map with the down-sampled coarse mask $M^{\prime }\in R^{H^{\prime }\times W^{\prime }\times 1}$ in the feature channel dimension to further incorporate the coarse annotation prior. To get the final fusion feature map $F^{\prime }\in R^{H^{\prime }\times W^{\prime }\times C^{\prime }}$, we adopt a simple 1-layer network with architecture: 1$\times$1 Conv2d$+$BatchNorm2d$+$ReLU.

Given the fused feature map, we want to learn representative and well-separated prototype vectors for the lesion region and the background based on the prototypical network. In previous research [17, 14, 22], the prototypical network condenses the masked object features in an image into a single or few prototypes. A relative simple coarse foreground lesion prototype can be calculated by mask average pooling, as follows:

where $(x,y)$ indexes the spatial locations and $1(\bullet )$ is an indicator function. In addition, the background prototype is computed by

where $p_{fg},p_{bg}\in R^{C^{\prime }}$.

Refinement is done using a non-parametric metric learning method [17]. For each pixel at location $(x,y)$ of the final fusion feature map $F^{\prime }$, we calculate the distance between its feature vector and the derived prototypes $P=\left\{p_{bg}{,p}_{fg}\right\}$. Then, we apply the softmax operation over the distances to get the probability map $P_c\in R^{H^{\prime }\times W^{\prime }\times 1}$ and $c\in \{bg,fg\}$. Formally, we have:

where $\alpha$ is the scaling factor fixed at 20.

We train our model end-to-end using the sum of dice loss $L_{dice}$ and binary cross-entropy loss $L_{bce}$ between the final probability map $P_{fg}$ and the well-annotated ground truth mask $M_{gt}$. That is: $L_{loss}=L_{dice}+L_{bce}$.

During testing, for each image patch $I$ and its corresponding coarse lesion annotation $M$, we obtain a corresponding foreground probability map $P_{fg}$. When mapping $P_{fg}$ back to the original image space (uncropped full image), some of them will overlap. For each pixel in the overlapping area, we choose the maximum probability of these probability maps as its value. In the end, thresholding is used to get the final refined mask.

There is no public available retinal lesion dataset with paired coarse annotation and pixel-level segmentation mask. To construct such paired dataset, we develop a simple coarse annotation generation method. Firstly, the coarse annotations are simulated from the well-annotated fine masks by applying the following chain of operations: smoothing, dilating, expanding, clustering the connected components using DBSCAN [2] and fitting ellipses to each cluster. Secondly, the fundus image is cropped around each ellipse in the corresponding coarse annotation. Finally, these cropped image patch and coarse annotation pairs are resized to fixed dimensions $H\times W$ for subsequent model training and testing.

We evaluate the proposed methods on publicly IDRiD and DDR datasets, and our real-world private dataset. IDRiD contains 81 fundus images (54 training images, 27 testing images) with pixel-level annotations for hard exudates (EX), hemorrhages (HE), microaneurysms (MA), and soft exudates (SE). Similarly, the testing part of DDR contains 225 fundus images with pixel-level annotations for EX, HE, MA, and SE. Our real-world private testing dataset collects 211 fundus images with pixel-level annotations for drusen (Drus) and pre-retinal hemorrhages (Prh) labeled by two experienced ophthalmologists. To train our annotation refinement network, we collect 32985 training patch pairs (EX:9957, HE:7752, MA:14387, SE:889) from the IDRiD training images using our coarse annotation generation algorithm. We also apply the mask generation algorithm to generate coarse mask for each testing image. To compare different refinement methods, we calculate the Intersection over Union (IoU) between the refined annotation and ground-truth mask.

We implement several non-prototypical mask refinement baseline models, taking in an image patch and a coarse mask as the input. In detail, we choose three widely used feature extraction backbones, Res18 [3], HRNet18 [23], and U-Net [13] attached with the coarse mask fusion module to perform feature extraction. The feature extraction process is identical to the one described in Sec.2.1.1 except for the feature backbone. After that, we attach a 1x1 Conv2d layer as a binary classifier to obtain a refined segmentation score map.

For the prototypical methods, we set the superpixel number $N_{sp}=20$. All training patch pairs are resized to 256 $\times$ 256 and augmented by RandomShiftScaleRotate, RandomBlur, and RandomBrightnessContrast. All models are implemented by PyTorch and trained from scratch using Adam optimizer with a batch size of 64 for 120 epochs. The initial learning rate is ${10}^{-4}$ and reduces according to ReduceLROnPlateau strategy.

We train the proposed coarse annotation refinement network using all four lesion types on IDRiD and evaluate the performance on the IDRiD testing set. As shown in Table 1, “Initial Coarse” denotes the IoU of actual lesion region versus coarse annotation region. Our prototypical method improves the initial coarse mask considerably. It also consistently obtains better refinement performance than the non-prototypical baselines on all four lesion classes in terms of IoU score, with average IoU score improving by more than 5.2%. This experiment demonstrates our advantages when training and testing images are from the same dataset.

We directly evaluate the performance on the DDR testing set and our real-world private dataset using models trained on the IDRiD dataset without further fine-tuning. As shown in Table 2, our method exceeds the U-Net baseline by 4.3% on both DDR and private datasets. For DDR, our superpixel weighted prototype performs better for all lesion types compared to the non-weighted prototype. Similarly, the weighted prototype is notably better than the non-weighted one on the private dataset, especially for the class Prh (56.3% $\to$ 62.6%). Overall, we see a general trend that our model can generalize well to new datasets or unseen classes.

Since ophthalmologists tend to draw a single rough ellipse to cover several unconnected lesion regions, we simulate the process by setting different reduction factors to the DBSCAN clustering algorithm. Actually, the number of the generated ellipses is the number of connected lesion regions divided by the reduction factor. In other words, with a higher reduction factor, the generated coarse mask will be more coarse. As shown in Fig. 3, although the refinement performance of all methods degrades as the reduction factor ranges from 1.0 to 2.0, our prototypical method has less degradation compared to the U-Net baseline. It implies our method is more robust against coarser annotations.

Fig. 4 presents some visualization of refinement results. Despite the vast variation in lesion scales, colors, and low contrast to surrounding regions, the first three rows show our proposed superpixel weighted prototype approach generates the most accurate lesion boundary. The last row shows a failure case when the coarse mask contains two distinct lesion classes, EX and Drus, at the same time.