Light-weight Retinal Layer Segmentation with Global Reasoning

Medical image segmentation has made rapid progress under the promotion of deep learning, especially in the fields of biological tissue recognition and lesion detection, which has sparked a large amount of research on medical image segmentation. The morphology of the optic disc and optic cup, as well as the cup-to-disc ratio, can be used to assess glaucoma, Wang et al. proposed an automatic segmentation approach based on CNNs to accurately segment the optic disc and optic cup from fundus images for glaucoma detection[16]. Precise segmentation of retinal vessels from fundus images is essential for intervention in numerous diseases and high-precision segmentation of retinal vessels still remains a challenging task, Li et al. proposed a dual-path progressive fusion network, named DPF-Net, which achieved better segmentation ability by effectively fusing global and local features[17]. Abdominal multi-organ image segmentation plays a crucial role in the diagnosis and treatment of many diseases. traditional methods of manual depiction are inefficient and subjective, therefore, Chen et al. proposed TransUnet based on CNN and Transformer block for abdominal multi-organ segmentation[18]. TransUnet was the first model to introduce the Transformer block into the U-shaped encoder-decoder structure and achieved very high accuracy in segmentation. There are also segmentation studies dedicated to lesion regions, such as Astaraki et al. who segmented different types of lung pathological regions to enable screening for lung diseases[19]. In addition, some lightweight segmentation methods for practical applications have also been proposed, such as SegFormer[20], A-net[21], EMV-Net[22].

The morphology of the retinal layer is associated with the development of many ophthalmic diseases, therefore, there are many studies on retinal layer segmentation based on CNN and OCT images. The success of ReLayNet opened the door to the application of CNNs for retinal layer segmentation[11], a typical U-shaped network based on an encoder-decoder structure, which achieved the best performance for retinal layer segmentation at that time. Since layer segmentation of OCT retinal images is prone to speckle noise, intensity inhomogeneity, Wang et al. proposed a boundary-aware U-Net for retinal layer segmentation by detecting accurate boundaries[23]. To solve the topological errors caused by not considering the order of retinal layers, Liu et al proposed a novel deep learning-based framework that employed the distance maps of layer surfaces to convert the layer segmentation task into a multitasking problem for classification and regression[24]. To ensure the continuity of the retinal layer boundary and obtain more accurate segmentation results, Hu et al proposed a coarse-to-fine retinal layer boundary segmentation method based on the embedded residual recurrent network and the graph search, it has achieved good performance in both qualitative and quantitative indicators[25]. In addition, considering the impact of different neurological diseases on the retinal layer, Gende et al presented a fully automatic approach for the retinal layer segmentation in multiple neurodegenerative disorder scenarios[26], the results indicate that the model is more robust. There are also studies aimed at practical application deployment, such as He et al. proposed a light-weight retinal layer segmentation approach that achieved good performance with a lower number of parameters[22].

Our proposed LightReSeg is a U-shape network based on an encoder-decoder structure as shown in Fig. 2. The network takes a retinal layer OCT image as input. The first $N$ scale feature maps ($F_{1}$, …, $F_{N-1}$, $F_{N}$) are extracted by the encoder which consists of a multi-scale feature extractor, respectively. In order to further deepen the global reasoning capability in the depth direction, the feature map $F_{N}$, which has to pass through a depth feature encoder: Transformer block[13]. The feature map $F_{TN}$ further encoded by the Transformer block has no change in shape and size, but it facilitates long-range global reasoning. To better preserve details at different scales, multi-scale features ($F_{1}$, …, $F_{N-1}$, $F_{N}$) are then merged with its corresponding layers in the decoder part via our proposed MAA module, which outputs the feature maps ($F_{M1}$, …, $F_{MN-1}$, $F_{MN}$) with the same shape size as its input. The $F_{TN}$ output from the Transformer block is merged with the $F_{MN}$ for feature fusion, and then after 2 fold size up-sampling, it continues to fuse with $F_{MN-1}$ for feature fusion and then follows a similar operation until the fusion of the up-sampled restored feature map with $F_{out}$ is completed, and after channel adjustment, the final retinal layer segmentation image is output. LightReSeg is designed for real-time clinical operations, where light-weight designs are employed where possible.

We present in full the multi-scale encoder in Section III-B followed by the proposed MAA module that serves for multi-scale feature fusion in Section III-C. Finally, we describe the design details for computation efficiency in Section III-D.

The multi-scale encoder extracts $N$ scale features to fully exploit the semantic information of feature maps at all scales, further, in order to make the features have better global reasoning capabilities, we incorporate a Transformer block[13], which is a module that further optimizes the $N$th scale feature map extracted from the deepest scale of the encoder. To ensure the dimensionality is consistent with the Transformer block port, we pre-process the feature map $F_{N}$ before being processed by the Transformer block. The feature map $F_{N}$ is first converted into a $2D$ patch sequence as

where $c$ is the number of channels in the $F_{N}$ feature map, $P$ is the size of the patch, and $H$ and $W$ are the height and width of the $F_{N}$ feature map.

We use a trainable linear projection to map the vector patch $x_{p}$ to a latent D-dimensional embedding space. To encode the spatial information of the patches, we learn specific position embeddings that are added to preserve the positional information in the patch embedding as

where $E\in\mathbb{R}^{\left(P^{2}\cdot C\right)\times D}$ represents the patch embedding projection, the $x_{0}$ is an additional learnable vector concatenated with the remaining vectors to integrate the information of all remaining vectors, the $\left[\cdot;\cdots;\cdot\right]$ is the concatenation operator, and $E_{pos}\in\mathbb{R}^{(Z+1)\times D}$ represents the position embedding[13].

Then, inputting $z_{0}$ into the Transformer Layer, as shown in Fig. 2, first goes through the first Layer Norm (LN) layer and then enters the Multi-head Self-attention (MSA) layer, here $MSA\left(LN\left(z_{0}\right)\right)+z_{0}$ forms the residual structure and get the $z_{0}^{\prime}$. Then, after passing through the second LN layer, it enters the Multi-Layer Perceptron (MLP) layer, here $MSA\left(LN\left(z_{0}^{\prime}\right)\right)+z_{0}^{\prime}$ once again forms the residual structure and get the $z_{1}$.

At this point, the first Transformer layer is finished. We have set three Transformer layers in our approach, so we have to cycle three times.

When the $z_{L}$ outputs from the Transformer block, we have to remove the $x_{{0}}^{{}^{\prime}}$ vector from the $z_{L}$, because only then we can reshape it to the same size of $F_{N}$.

To better preserve the semantic information at each encoder scale, the MAA module is designed in the skip connection part of the encoder and decoder by a large convolution kernel and channel attention mechanism. It’s also set with multi-scale feature fusion, and it’s based on the ideas of asymmetric convolution[15], multi-scale feature extraction[27, 28] and channel attention mechanism[29].

The workflow of the MAA module is illustrated in Fig. 3, it begins with the input of a feature map $F_{0}$ extracted from the feature encoder, assuming a shape size of ($C$, $H_{0}$, $W_{0}$). $F_{0}$ enters MAA and goes through a $5\times 5$ convolution kernel to expand the perceptual field, and then has to enter a multi-scale convolution to extract information at different scales (Eq. 4, Eq. 5). Each scale goes through a channel attention (CA) module to increase the weight of the feature map on the channel dimension (Eq. 6). As shown in the channel attention module in Fig. 3, the feature $F_{0i}^{{}^{\prime}}\in\mathbb{R}^{C\times H\times W}$ reshape to $F_{0i\_cn}^{{}^{\prime}}\in\mathbb{R}^{C\times N}$, where $N=H\times W$, $\in$ represents the shape of the feature map. After that we perform a matrix multiplication between $F_{0i\_cn}^{{}^{\prime}}$ and the transpose of it, get the $F_{0i\_cc}^{{}^{\prime}}\in\mathbb{R}^{C\times C}$. Eventually, we apply the softmax layer to get the channel attention map $A\in\mathbb{R}^{C\times C}$

(Eq. 7, Eq. 8). Here the matrix $A\in\mathbb{R}^{C\times C}$ contains information about the weights between channels, e.g. $\left(A\in\mathbb{R}^{C\times C}\right)_{ij}$ represents the impact of $i^{th}$ channel on $j^{th}$ channel. In addition, we multiply the matrix between the transpose of $A\in\mathbb{R}^{C\times C}$ and

$F_{0i\_cn}^{{}^{\prime}}$, and reshape the result to $\mathbb{R}^{C\times H\times W}$. Then we multiply the result by a scale parameter $\alpha$ and perform an element-wise sum operation with $F_{0i}^{{}^{\prime}}$ to obtain the final output $F_{0i}^{{}^{\prime\prime}}$ (Eq. 9 and Eq. 10). After the CA module, all feature maps are summed and convolved by $1\times 1$ kernel to generate a weighted feature map, and we use the weighted feature map to re-weight $F_{0}$ by element-wise multiplying ($\bigotimes$) it with the input $F_{0}$ (Eq. 11 and Eq. 12).

At last, we obtain the $F_{M0}$. According to the above procedure, we can input the extracted four scales feature maps $F_{1}$, …, $F_{N-1}$, $F_{N}$ into MAA module to get $F_{M1}$, …, $F_{MN-1}$, $F_{MN}$, respectively, as shown in Fig. 2.

Instead, the ASPP employs a series of dilated convolution layers with different rates to obtain larger receptive fields, however, dilated convolutions may lose local detail information. In contrast, the MAA module utilizes asymmetric convolutions with large kernels, which allows capturing multi-scale features without sacrificing too much detail. Furthermore, before the 1x1 convolution compresses the channels in the MAA module, a channel attention mechanism is adopted, enhancing the module’s focus on the channel dimension and better preserving important features.

The basic principle of our design of LightReSeg is to improve the overall performance of the model with as few parameters as possible, so we can meet the computational efficiency of deployments for real-time clinical applications. Thus, we make some light-weight designs for part of the structure of the model. For encoder extraction of multi-scale feature maps, we use depthwise separable convolutions (DS-Conv) for our down-sampling[14]. DS-Conv considers the channels and spatial regions of the feature map separately, takes different convolution kernels for different channels, and then uses point convolution to aggregate the channel information, achieving a richer feature representation with fewer parameters to learn. We find that using standard convolution for down-sampling is better than pooling, but convolution increases the number of parameters, so we use DS-Conv to ensure the segmentation performance while decreasing the burden of the number of parameters. In addition, as mentioned in Sec. III-C, multi-scale asymmetric convolution is used in our MAA module. Initially, our idea is to perform multi-scale extraction of feature information like ASPP [28], but we find that using standard convolution would cause a greater burden on the number of parameters, since the MAA module is supposed to perform extraction of the $N$ scale feature maps, which would certainly further increase the number of parameters. In order to reduce the number of parameters in the MAA module, we use asymmetric convolution instead of standard convolution, shown in Fig. 3, and the overall performance of LightReSeg is nearly the same. In particular, the Transformer block has a significant proportion of the number of parameters in LightReSeg. We have optimized the internal parameters of the Transformer block structure, such as the setting of heads=8 for multiple heads in MSA and dim_head=64 for hidden linear layers in MSA[13], etc. All these designs to minimize the number of parameters are made under the precondition of ensuring the performance of LightReSeg.

We evaluate the segmentation performance quantitatively using the Dice similarity coefficient (DSC), mean Intersection over Union (mIoU), pixel accuracy (PA), and mean pixel accuracy (mPA). They can be computed as:

and

where $TP$, $FP$, $TN$, and $FN$ represent True Positives, False Positives, True Negatives, and False Negatives respectively. $k$ represents the number of categories. In addition, we count the number of parameters for compared approaches to represent their computation complexity.

LightReSeg is implemented based on the PyTorch and trained with the Adam optimizer with the cross-entropy loss. The initial learning rate is set to 0.001 which is then gradually halved every 40 epochs. Data augmentation is applied on all three datasets, including horizontal flipping with probability P=0.5, random center rotation within plus or minus 20 degrees with probability P=0.5, median and motion blur processing with P=0.5 probability, random Gaussian noise addition, as well as random brightness and contrast with P=0.5 probability. All experiments are conducted on an NVIDIA GeForce RTX 3090 Graphics card. The code of the proposed LightReSeg could be found at: https://github.com/Medical-Image-Analysis/LightReSeg.

We compare LightReSeg with the state-of-the-art approaches including ReLaynet[11], EMV-Net[22], Attention_Unet[12], BiSeNet[32], DFANet[31], OS_MGU[30], U-net [10] and TransUnet[18]. We report the results on the above-mentioned three datasets.

Although LightReSeg boasts only 3.3M trainable parameters and delivers the optimal retinal layer segmentation outcomes, in real-world applications, model parameter size is not the sole determinant of algorithm efficacy. Inference speed emerges as a crucial metric for assessing algorithmic effectiveness. As evidenced by Tab. VII, our proposed method exhibits inference speeds of 0.11$s$, 0.27$s$, and 0.07$s$ on the three datasets, respectively. While these figures already signify remarkable efficiency, there remains scope for further enhancement. For example, although the segmentation accuracy of ReLayNet is far inferior to our approach, it must be admitted that its inference efficiency is slightly higher than ours.

Moreover, in the context of practical clinical applications, constraints related to dataset limitations preclude the inclusion of all device data types. Imagery obtained through disparate devices may exhibit domain discrepancies, inevitably resulting in false positive errors in the segmentation output for unfamiliar new devices. Consequently, apart from enlarging the dataset, enhancing the model’s domain generalization competency is a subject worthy of further investigation. The supplementary segmentation result evaluation system will also augment ophthalmologists’ confidence in the segmentation outcomes.