The classification method based on convolutional neural network can improve the performance of the pattern recognition system by automatically extracting and self-learning features. At present, the conventional image classification network, including the network method designed in this paper, directly uses the common depth convolution network to directly extract and classify features. Therefore, the core of optimizing the experimental results is how to better extract the features in the image, that is, to extract more advanced and richer features. The advantage of Transformer is to capture the global context information in the way of attention, so as to establish a long-distance dependence on the target, so as to extract more powerful features.
Research Article
ResTr: A Joint Framework for Retinal Vein Occlusion Image Classification
https://doi.org/10.21203/rs.3.rs-2225149/v1
This work is licensed under a CC BY 4.0 License
Version 1
posted
You are reading this latest preprint version
Over the past few years, transformer has demonstrated superior performance in the field of natural language processing[1]. Different from the convolutional neural network (CNN), which focuses on the local receptive field of convolutional layer, transformer mainly extracts the intrinsic features of text data based on self-attention mechanism. Due to its powerful feature extraction and representation capabilities, researchers have been trying to apply transformer to computer vision tasks.
The good news is that transformer has also made significant breakthroughs in the field of computer vision during the last two years. A large number of transformer-based methods have been applied to image processing tasks, such as DETR[2] for object detection, ViT[3] and DeiT[4] for image classification, and SETR[5] for semantic segmentation. However, compared to natural images, medical images are usually limited in number due to the high cost of data collection and annotation. And they are larger in size, with more noise and redundant information. Therefore, how to optimize the performance of transformer-based models in medical image tasks is a problem that needs to be addressed continuously.
ResNet and its variant structures are currently the basic building blocks of most methods proposed for image classification. The convolution kernels of these networks basically focus on only a local subset of pixels in the whole image and then strictly follow the process of refining global features from local features. But as thus, the networks will lack the ability to model the long-range correlations present in the image. Nevertheless, in medical images, such as the retinal images that are the focus of this paper, there are many explicit sequences that contain important long-range dependencies and semantic information. As shown in Figure, most visual representations in retinal images are ordered, and disruption of these sequences would significantly reduce the accuracy of the model classification. The previously mentioned transformer is an encoder that handles serial relationships, which enables global or remote interaction between all embedded entries. In Paper [6], Google examined some key layers of ResNet and ViT using Centered Kernel Alignment (CKA) and demonstrated that ViT retains more spatial information than ResNet. This shows that transformer has a powerful ability to model sequence correlation and long-range correlation of semantic information. And it is our motivation for applying transformer to retinal image classification.
As described in other articles[7–17], convolutional neural network has more or less many shortcomings.
For instance, Jeya et al. proposed a gated axial-attention model, which extended the existing architectures by introducing an additional control mechanism in the self-attention module. Shuang Yu et al. proposed a multiple instance learning (MIL) based ‘MIL head’, which can be conveniently attached to the ViT in a plug-and-play manner. Hu Cao et al. used hierarchical Swin Transformer with shifted windows as the encoder to extract context features. These works basically revolve around pure transformer structures, but transformer-based models only work well when trained on large-scale datasets. In addition, due to the small size of the retinal image dataset, it would be a problem that the network lacks sufficient information to establish relationships between low-level semantic features.
To overcome the above limitations, we designed a joint framework that can combine the transformer with the CNN structure. Two major contributions are achieved with this paper:
1) We propose a joint model ResTr for retinal image classification. we introduce transformer to explore the long-range dependence of explicit sequences of images and employ CNN to compensate for the loss of local structural information.
2) In the ResTr method, we design a joint loss function to optimize its dual flow and increase the intra-group aggregation.
In this paper, we propose ResTR, and its framework is provided in Figure. As shown in Figure, ResTr consists of two different branches to extract semantic features and structural features, respectively. Furthermore, we design a joint loss function to optimize the end-to-end model. As shown in the figure below:
2.1 CNN Branch
As a CNN model widely used in image classification tasks, ResNet is used to extract local features. We used the core of ResNet's architecture, the residual structure other than the fully connected layer. The residual structure can solve the gradient and performance degradation problems while maintaining network complexity and depth, and achieve very significant classification results. Unlike other traditional CNNs, ResNet uses stride = 2 convolution for downsampling and replaces the fully connected layer with a global average pool layer. Feature extraction is then performed by stacking 3X3 convolutions, and training is accelerated using Batch Normalization (i.e., dropout is discarded). Figure 3 shows the network structure of ResNet, which is conventionally divided into 5 parts.
2.2 Transformer Branch
We use the transformer to extract global features thanks to its ability to model the long-range dependencies between the input sequence elements. Unlike CNN branches, the input image 
needs to be divided into separate patches with size of P × P[ViT]. An image can be obtained with 16 patches, which are further flattened into 1D format and then embedded into the D dimension by linear layers, which are further flattened into 1D format and then get embedded via a linear layer into D dimensions. For the patch embedding the position information of the image sequence is added using the standard learnable 1D position embedding. The resulting vector sequence is then fed into the Transformer encoder. More specifically, the encoder consists of alternating Multihead Self-Attention (MSA) and MLP block layers. LayerNorm (LN) is applied before each block and residual concatenation is applied after each block.
2.3 Self-Attention Overview
The core of the Transformer structure is the SA mechanism, which is similar to the idea of attention. In the image classification task of this paper, SA can be understood as computing the correlation between each pixel point. The specific computation process is following: an input feature map with height H, weight W and channels C. The output of a self-attention layer is computed using the following equation:
actual process of getting such affinities, the calculation is very expensive. Inspired by axial attention[Axial-DeepLab], self-attention is decomposed into two self-attention modules. The first module performs self-attention on the feature map height axis and the second one operates on the width axis. Axial attention can effectively capture non-local contextual information and is computationally more efficient. It can also effectively embed location information and capture long-range dependencies between feature mappings. After that, the updated self-attention mechanism along with width axis can be written as:
Eq. 2 describes the axial attention applied along the height axis of the tensor, and a similar formulation is also used to apply axial attention along the width axis after. Figure 2 illustrates the calculation process of axial attention along two directions.
Dataset Details
The main manifestations of RVO fundus are tortuous dilatation of the affected veins, flame-like haemorrhage along the retinal veins[7]. Each eye of the patient is treated as an independent classification sample. Multicolor(MC) image is one of the new technologies based on confocal scanning laser ophthalmoscope (cSLO), which can obtain multiple modal images with different wavelengths. We collected the results of the MC image examination from 29 patients as an internal dataset with a total of 220 images (some cases with missing images). The number was increased to 1320 by image enhancement strategies, including random cropping, rotation, horizontal flipping and color jitting. The ophthalmologist manually labelled them as RVO or not. These labels are treated as the ground truth for validating our algorithm. As shown in the figure below:
Experimental Setting
We implement the proposed framework and perform experiments using the PyTorch library. All experiments are performed on a GPU cluster with four NVIDIA GeForce RTX 3090 GPUs and each with 24 GB of memory. The optimization process is run for 50 epochs. We set the learning rate as 0.0001, and the batch size is set as 12. All methods are optimized by the Adam optimizer. The cross-entropy function is selected as the classification loss function.
One of the research difficulties is that Transformer mechanism is a relatively new image processing technology introduced from the field of natural language processing, which is not as mature as the application of convolutional neural network in the field of computer vision. Moreover, the reason why Transformer has strong ability to learn long-term semantic information is that it does not assume to start from local information, but can get global information from the beginning, so it will be more difficult to learn. In view of these difficulties, the improved network structure in this paper requires a lot of experiments to constantly adjust the parameters to optimize the classification results.
Less pathological sample data is a major difficulty in medical image analysis task, which is also inevitable in this study. Although the number of data sets can be expanded through data enhancement, in order to avoid over fitting of classification model training, it is also necessary to balance data expansion and diversity. In addition, in order to improve the generalization ability of the model, we need to continue to find large public datasets for comparative experiments.
- Vaswani A, Shazeer N, Parmar N et al (2017) Attention is all you need[J].Advances in neural information processing systems,30
- Carion N, Massa F, Synnaeve G et al (2020) End-to-end object detection with transformers[C]//European conference on computer vision. Springer, Cham, : 213–229
- Dosovitskiy A, Beyer L, Kolesnikov A et al An image is worth 16x16 words: Transformers for image recognition at scale[J]. arXiv preprint arXiv:2010.11929, 2020.
- Touvron H, Cord M, Douze M et al (2021) Training data-efficient image transformers & distillation through attention[C]//International Conference on Machine Learning. PMLR, : 10347–10357
- Zheng S, Lu J, Zhao H et al (2021) Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. : 6881–6890
- Raghu M, Unterthiner T, Kornblith S et al (2021) Do vision transformers see like convolutional neural networks?[J]. Adv Neural Inf Process Syst 34:12116–12128
- Usman M, Fraz MM, Barman SA (2017) Computer vision techniques applied for diagnostic analysis of retinal OCT images: a review[J]. Arch Comput Methods Eng 24(3):449–465
- Singh A, Dutta MK, ParthaSarathi M et al (2016) Image processing based automatic diagnosis of glaucoma using wavelet features of segmented optic disc from fundus image[J]. Comput Methods Programs Biomed 124:108–120
- Fu H, Cheng J, Xu Y et al (2018) Disc-aware ensemble network for glaucoma screening from fundus image[J]. IEEE Trans Med Imaging 37(11):2493–2501
- Singh A, Dutta MK, ParthaSarathi M et al (2016) Image processing based automatic diagnosis of glaucoma using wavelet features of segmented optic disc from fundus image[J]. Comput Methods Programs Biomed 124:108–120
- Li Z, He Y, Keel S et al (2018) Efficacy of a deep learning system for detecting glaucomatous optic neuropathy based on color fundus photographs[J]. Ophthalmology 125(8):1199–1206
- Li A, Cheng J, Wong DWK et al (2016) Integrating holistic and local deep features for glaucoma classification[C]// 2016 38th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC). IEEE, :1328–1331
- Szegedy C, Vanhoucke V, Ioffe S et al (2016) Rethinking the inception architecture for computer vision[C]// Proceedings of the IEEE conference on computer vision and pattern recognition. : 2818–2826
- Karri SPK, Chakraborty D, Chatterjee J (2017) Transfer learning based classification of optical coherence tomography images with diabetic macular edema and dry age-related macular degeneration[J]. Biomedical Opt express 8(2):579–592
- Wang J, Wang Z, Li F et al (2019) Joint retina segmentation and classification for early glaucoma diagnosis[J]. Biomedical Opt express 10(5):2639–2656
- Acharya UR, Chua CK, Ng EYK et al (2008) Application of Higher Order Spectra for the Identification of Diabetes Retinopathy Stages[J]. J Med Syst 32(6):481–488
- Du N, Li Y (2013) Automated identification of diabetic retinopathy stages using support vector machine[C]// Control Conference (CCC), 2013 32nd Chinese. IEEE, : 3882–3886
No competing interests reported.