Transformer-Based Disease Identification for Small-Scale Imbalanced Capsule Endoscopy Dataset

Abstract

Vision Transformer (ViT) is emerging as a new leader in computer vision with its outstanding performance in many tasks (e.g., ImageNet-22k, JFT-300M). However, the success of ViT relies on pretraining on large datasets. It is difficult for us to use ViT to train from scratch on a small-scale imbalanced capsule endoscopic image dataset. This paper adopts a Transformer neural network with a spatial pooling configuration. Transfomer’s self-attention mechanism enables it to capture long-range information effectively, and the exploration of ViT spatial structure by pooling can further improve the performance of ViT on our small-scale capsule endoscopy dataset. We trained from scratch on two publicly available datasets for capsule endoscopy disease classification, obtained 79.15% accuracy on the multi-classification task of the Kvasir-Capsule dataset, and 98.63% accuracy on the binary classification task of the Red Lesion Endoscopy dataset.

1. Introduction

Gastrointestinal (GI) cancer is the second most deadly cancer worldwide [1], accounting for 35% of cancer-related deaths [2]. Endoscopy and radiological techniques can provide physicians with visual signals of the GI environment to diagnose. Compared to endoscopic techniques, radiological techniques produce inferior images with less discernible features. However, conventional endoscopes can only be inserted into the stomach and large intestine, but the small bowel structure is often overlapping and has no fixed shape, making it difficult for conventional endoscopes to penetrate. It was challenging to examine the small bowel until wireless capsule endoscopy (WCE) was developed [3,4]. Capsule endoscopy can easily access the small intestine and provide direct, clear video signals for physicians to observe possible lesions in the small bowel, such as ulcers, erosions, angiodysplasias, lymphoma, etc. [5,6].

WCE is a particular endoscope with imaging equipment, batteries, light source, and a signal-transmitting component. It is so tiny that a patient can directly swallow it. The WCE can move passively within a patient’s gastrointestinal environment or selectively move and examine the patient’s lesions under external control (e.g., magnetic control). Then, WCE will output a large amount of video information for the physicians’ examination to determine the disease and follow-up treatment plan. However, the frame-by-frame analysis of the WCE output video is quite tedious. The examiner needs to be knowledgeable and experienced in small bowel disease and be focused at all times to identify the many different lesions [7,8]. These difficulties lead to high missing and misdiagnosis rates when the physicians manually deal with the WCE video information [9,10]. The WCE system overview can be found in Figure 1. A particular signal receiver will receive the image information transmitted by the capsule endoscopy inside the patient. Subsequently, a system based on artificial intelligence (AI) algorithms will analyze the acquired images in real time, thereby assisting physicians in making a diagnosis as quickly as possible.

With the rapid development of automation and visual analysis technology, deep learning (DL) has evolved and significantly grown in the last decade. DL methodologies have been widely used in various fields [11,12,13,14,15,16,17,18]. More and more DL algorithms are applied in automatic WCE diagnosis and analysis, and they have achieved excellent results. DL algorithms can bring very high diagnostic accuracy while helping to reduce the labor of physicians [19,20]. Meanwhile, some researchers apply DL to depth and motion estimation of WCE [21,22]. DL methodologies have shown great potential in WCE data, and will provide excellent help and improvement to future WCE technology.

The transformer is a recent emerging DL architecture [23]. Its architecture consists entirely of the self-attention mechanism and achieves better performance than Convolutional Neural Network (CNN). A. Dosovitskiy et al. [24] introduced the transformer into the task of image classification. Vision Transformer (ViT) showed extraordinary potential in image processing. However, the training of transformers often requires a large amount of data, and its excellent performance comes from pretraining based on large datasets. Even ImageNet [25] sometimes fails to unleash the full performance of ViT [24]. This reliance on pretraining is explained as locality inductive bias [26], requiring more data to learn features and visual representations. However, the datasets collected by capsule endoscopy usually have the following three major problems:

1.

The amount of data is small;

2.

The images of the dataset lack distinguishable features;

3.

The data representing different diseases are not balanced. Common diseases may be represented in a large number of images, while for some rare diseases, there are only a few images for training.

This paper adopts a transformer network that imitates the existing pooling architecture. It can be trained from scratch and applied to the small and imbalanced capsule endoscopy dataset. We conduct experiments on the Kvasir-Capsule dataset [27] and Red Lesion Endoscopy (RLE) dataset [28], obtaining results that outperform existing baselines.

The rest of this paper is organized as follows: Section 2 and Section 3 will introduce the related work and adopted methodology, respectively. Section 4 shows our experiments and results, followed by our discussion in Section 5. Finally, we conclude in Section 6.

2. Related Work

2.1. WCE Pathological Diagnosis

After the excogitating of WCE, pathological diagnosis methods regarding WCE have been seeking surmounts. In the early stage, researchers worked on developing publicly available WCE datasets. Several datasets, such as Kvasir-Capsule [27], RLE [28], KID [29], GIANA [30], have been published for researchers to reference and verification. Machine Learning methods have been widely developed for disease classification and detection in WCE images, such as support vector machines (SVM) [31] and single-shot multiBox detector (SSD) [32]. H. L. Gjestang et al. [33] presented a semi-supervised teacher–student framework with better performance than traditional supervised learning-based models—they utilized unlabeled samples for self-adjustment. As a result of the current popularity of self-attention mechanisms in computer vision, P. Muruganantham et al. [34] proposed a dual branch CNN model using self-attention mechanisms to bring out accuracy for WCE to classification and lesion localization. Traditional supervised DL methods require extensive samples to train and verify. However, obtaining large and specific lesion category samples from realistic clinic practice is challenging. R. Khadka et al. [35] proposed a new implicit model-agnostic meta-learning (iMAML) algorithm, which only requires small training samples, but performs high accuracy on unseen datasets.

Nevertheless, based on our discussion in Section 1, the most critical problems for the WCE datasets are the lack of distinguishable features, imbalanced classes, and the small amount of data. The commonly used convolutional calculation has space constraints, and the network’s attention area is relatively limited. ViT can learn the most appropriate inductive bias according to the task objective and the layer’s position in the network. Furthermore, the multi-head mechanism ensures that ViT can focus on multiple discriminative parts. These advantages allow ViT to achieve better results on WCE tasks.

2.2. Vit Training from Scratch

Several approaches have been proposed to overcome the problem that ViT must rely on large datasets for pretraining. Data-efficient Image Transformer (DeiT) [36] introduced a ViT-based teacher–student framework that allows student models to learn from teacher models through an attention mechanism. This ViT-based knowledge distillation significantly reduced training time and computational resource consumption. Tokens-To-Token (T2T) [37] learned the deep-narrow structure of CNNs and developed a layerwise Tokens-to-Token transformation methodology to learn the local structure around tokens. Convolutional vision Transformer (CvT) [38] attempted to bring the CNN’s ability to extract low-level information to ViT. It used Convolutional Token Embedding and mimicked CNN’s process of learning features, which combined the benefits of both CNN and ViT. Swin Transformer [39] merged image patches to construct a hierarchical Transformer, and used attention computed in a non-overlapping local window. Class-Attention in Image Transformer (CaiT) [40] proposed LayerScale to make the convergence of Deep ViT easier and used class-attention layers to make the processing of class tokens more efficient. Y. Liu et al. [41] utilized a self-supervised architecture for learning spatial information in a single image. It can work in conjunction with supervised learning of ViT, allowing ViT to be more robust and accurate when trained on small-size datasets. S. H. Lee et al. [42] proposed two generic add-on modules, Shifted Patch Tokenization and Locality Self-Attention, to solve the problem of the locality inductive bias.

However, despite the tremendous progress in training from scratch on natural datasets, ViT, as a state-of-the-art technique, has not been popularized in processing WCE data. Compared with the tiny dataset commonly used in general computer vision, the WCE dataset has fewer data and is highly imbalanced. Thus, it is not easy to directly port existing methods for handling small-size datasets, hindering the further promotion and development of ViT on the WCE dataset.

3. Methodology

3.1. Vision Transformer

ViT combines knowledge from computer vision and natural language processing. The original ViT [24] consists of patch embedding, position embedding, transformer encoder, class token, and classification head. The patch embedding divides the original image into $16\times 16$ patches and flattens the patches into a $1\times n$ sequence. Position embedding provides the model with the location information of the patch so the model does not need to learn patch stitching based on the semantics of the patches, saving the learning cost. The transformer encoder is used to process the flattened patches. It has LayerNorm (LN), multi-head attention, LN, multi-layer perceptron (MLP) in order. The transformer encoder is established based on the multi-head self-attention mechanism, which assigns different weights to the input information to aggregate information. ViT builds the network by stacking multiple transformer encoder blocks. After being processed by the transformer encoder, the patches are still $1\times n$ matrices. The class token is a separate vector used to aggregate the image feature of the corresponding processed patch. When connecting the class token to the linear classifier, the classification-predicted results can be obtained. Figure 2A shows the typical architecture of ViT.

3.2. Pooling in CNN

Pooling is a commonly used module in CNNs, which is also known as downsampling [43,44,45]. It is usually used after the convolutional layer to reduce the feature map dimension. It can effectively reduce the network parameters, prevent overfitting, improve model robustness, and reduce information redundancy.

Max Pooling and Average Pooling are the two most commonly used pooling methods. Max Pooling selects the maximum pixel value in the image as the value after Pooling. The gradient is backpropagated through the maximum value of the forward propagation process, and the gradient at other locations is 0. Max Pooling is more concerned with the local features in the image. Average Pooling involves selecting the average value in the image region as the value after pooling, which focuses more on the global features of the image. Mix Pooling is a combination of Max Pooling and Average Pooling, and it also has the complementary advantage of the two methods. Researchers utilize Concat operation and Add operation to combine Max Pooling and Average Pooling. ResNet [43] uses the Add pooling operation, which adds the corresponding feature maps and then performs convolution operations. The corresponding feature maps share a convolution kernel. DenseNet [44] adopts Concat pooling operation. Each channel corresponds to one convolutional kernel with higher computational consumption. All these network architectures show a spatially reduced dimensionality design.

3.3. Pooling in ViT

In the original ViT, the spatial dimension of each layer remains unchanged, and each layer has the same number of tokens. Inspired by [46], the pooling operation in ViT can effectively improve the model’s generalization ability. We can learn from the network architecture in CNN to obtain a downsampling ViT architecture in the spatial dimension.

However, the original ViT-processed data are presented as 2D matrix data, so we cannot directly apply pooling operation to them. B. Heo et al. [46] present pooling layers in ViT. They first reshape the 2D matrix back to 3D at the beginning of the pooling layer. Thus, the depthwise convolution can be performed here to conduct the common pooling operation in CNN, with the variation of the space dimension and number of channels. Then, the 3D data will be reshaped again into the 2D matrix, which can be added back into the ViT module. Moreover, the class token is spatially aligned with ViT’s pooling module using fully connected layer extensions. Figure 2B shows the detailed spatial size of our network architecture. We follow the same setup in the original ViT and only modify the spatial dimensions in the transformer encoder. We design the whole structure in the style of 3-step pooling so that the spatial dimension will be reduced three times. The final prediction result will be obtained through the classifier head after three consecutive downsamplings of the processed data.

4. Experiment Validation

4.1. Datasets

The experiments were conducted based on the Kvasir-Capsule dataset [27] and RLE dataset [28]. Their image examples and data distribution are attached in Figure A1 and Figure A2 and

Table A1. Data distribution of Kvasir-Capsule dataset [27]. The eleven categories of this dataset show strong imbalances.

Category	Number of Images
Normal Clean Mucosa	34,338
Reduced Mucosal View	2906
Lymphangiectasia	592
Erythema	159
Angiectasia	866
Blood—fresh	446
Blood—Hematin	12
Erosion	506
Ulcer	854
Polyp	55
Foreign Body	776

and

Table A2. Data distribution of RLE dataset [28].

Category	Number of Images
Normal	2160
Bleeding	1125

. In the Kvasir-Capsule dataset, we remove the three classes from the Anatomy category, and retain the Luminal findings category. Thus, it has 11 classes, including 10 different categories of diseases (Angiectasia, Blood-fresh, Blood-hematin, Erosion, Erythema, Foreign Body, Pylorus, Lymphangiectasia, Ulcer, Reduced Mucosal View) and Normal Clean Mucosa images. The image resolution is 336 × 336. The data distribution of the Kvasir-Capsule dataset is very imbalanced—the largest class has 34,338 images, and the smallest class has only 12 images. The RLE dataset was originally for tissue hemorrhage segmentation in capsule endoscopy images. Its image resolution is 320 × 320. We rearranged the dataset into two classes, normal and bleeding images, and performed a binary classification task on this dataset.

Meanwhile, to ensure the robustness of prediction and prevent over-fitting, data augmentation strategies such as random rotation, flipping, and random noise are used to expand the datasets and improve their diversity and generalization ability. The training and testing sets are divided by the ratio of 4:1.

4.2. Model Evaluation

The neural networks were completed based on the Python PyTorch framework. The training and validation were conducted on a server with Ubuntu 18.04 LTS operating system, 2 NVIDIA GeForce RTX 3090 GPUs, and an AMD Ryzen™ 9 5950X CPU. We set the batch size to 128, and all models were trained for 100 epochs. The learning rate is set to 0.001 with the Adam optimizer. The comparison algorithms can be split into CNNs and ViTs as follows:

• CNNs: ResNet32 [43], ResNet56 [43], VGGNet19 [47], WRN16-2 [48], WRN40-2 [48], MobileNetV2 [49], ShuffleNetV1 [50], ShuffleNetV2 [51];
• ViTs: Swin-Transformer [39], CaiT [40].

We implemented all the experiments in our environment.

Table 1. The results of comparison experiments.

Models	Kvasir-Capsule	RLE	Inference Time/Per Image
ResNet32 [43]	74.49%	95.22%	0.41 ms
ResNet56 [43]	74.76%	95.77%	0.42 ms
VGGNet19 [47]	74.77%	94.30%	0.41 ms
WRN16_2 [48]	74.52%	93.01%	0.41 ms
WRN40_2 [48]	74.50%	93.38%	0.42 ms
MobileNetV2 [49]	74.91%	94.67%	0.41 ms
ShuffleNetV1 [50]	75.62%	95.40%	0.42 ms
ShuffleNetV2 [51]	75.33%	94.49%	0.41 ms
Swin [39]	77.88%	98.33%	0.47 ms
CaiT [40]	78.63%	97.26%	0.48 ms
Ours	79.15%	98.63%	0.47 ms

and Figure 3 show the classification accuracy and comparison. Firstly, among all CNN algorithms, ShuffleNetV1 achieves very excellent results on both datasets. ResNet56 achieves the highest accuracy among all CNNs on the RLE dataset, a binary classification task, but slightly lower than ShuffleNetV1 on the Kvasir-Capsule dataset, a multiclassification task. Secondly, in ViT algorithms, both CaiT and Swin Transformer show superior performance to CNN on both datasets. However, CaiT features a performance that exceeds Swin Transformer on the Kvasir-Capsule dataset, while Swin Transformer defeats CaiT on the RLE dataset. Lastly, the method we employed outperforms both CaiT and Swin Transformer quite clearly on both WCE datasets, with the accuracy of 79.15% on the Kvasir-Capsule dataset and 98.63% on the RLE dataset. We also present the inference time of different methodologies for a single image. ViT-based methods are slightly slower, but in real-time diagnostic applications, a difference of 0.06 ms (i.e., $6\times {10}^{-5}$ s) does not have a substantial impact.

5. Discussion

This section compares our methodology with other related and recently published WCE classification solutions. The associated solutions for WCE classification are listed in

Table 2. Comparison of the classification results of this study and other studies.

Authors	Methodology	Dataset	Accuracy
M. Sharif  et al. [52]	Geometric Features, CNN	10 WCE videos	99.1%
F. Rustam  et al. [53]	MobileNet, BIR	Bleeding Detection, Binary Classification	99.3%
X. Zhao  et al. [54]	CNN Backbone, LSTM, Transformer	113 WCE videos	93.0%
H. L. Gjestang  et al. [33]	Teacher–student Framework	Kvasir-Capsule [27]	69.5%
H. L. Gjestang  et al. [33]	Teacher–student Framework	HyperKvasir [55]	89.3%
P. Muruganantham  et al. [34]	Self-attention CNN	Processed Kvasir-Capsule [27]	95.4%
P. Muruganantham  et al. [34]	Self-attention CNN	Processed RLE [28]	95.1%
S. Biradher  et al. [56]	CNN	Processed RLE	98.5%
D. Bajhaiya  et al. [57]	DenseNet121	Ulcer Classification [27]	99.9%
N. Goel  et al. [58]	Dilated Input Context Retention CNN	8 WCE videos	96.0%
N. Goel  et al. [58]	Dilated Input Context Retention CNN	KID [29]	93.0%
A. Srivastava et al. [59]	FocalConvNet	Processed Kvasir-Capsule [27]	63.7%

.

Most of the research on WCE classification is based on CNN and its variants. In addition to exploring the existing CNN architectures [53,56,57], F. Rustam et al. [53] and N. Goel et al. [58] tried to improve the spatial architecture of the original CNN. A. Srivastavaet et al. [59] used the Focal Modulation Guided method to integrate global context information to CNN. M. Sharif et al. [52], P. Muruganantham et al. [34], X. Zhao et al. [54] added external information or components to CNN to improve the prediction ability. H. L. Gjestang et al. [33] explored the semi-supervised teacher–student framework to overcome the difficulties and deficiencies of annotation in medical data. Most of the above methods are based on CNN attempts, while X. Zhao et al. [54] introduced a transformer to process the WCE videos. However, they still used CNN as the feature extractor while utilizing the transformer to process temporal features and capture the interframe dependencies in videos.

Our methodology is based on a purely ViT network without the convolutional operations from the CNN. We employ spatial downsampling to integrate the spatial structure of ViT.

M. Raghu et al. [60] dissects the difference between CNN and ViT in terms of receptive field in detail. Usually, CNNs have only a fixed-size kernel (size 3 or 7), so CNNs cannot obtain global perception in the initial layer. Instead, the CNN obtains the local receptive field at initial layers and then gradually expands the receptive field by repeatedly “convolving” the information around the kernel layer by layer. Therefore, researchers must keep stacking convolutional layers to obtain a sizable receptive field via CNN. In contrast, based on the self-attentive mechanism of ViT, the receptive field (i.e., the mean attention distance of the self-attention head) of ViT in the initial layers has the local receptive field like CNN, but also has the global field of view. ViT’s global and long-range perception will ensure that when ViT makes the decision, the features ViT has used already cover sufficient information from the original image. Obviously, this also brings better performance to our disease classification for WCE. From Table 1, we know that the CaiT [40] and Swin Transformer [39] already outperform all CNN backbones. We also adopted spatial pooling to integrate the spatial structure of ViT. We can think of this as ViT learning from CNN to build the network structure, and then explore the potential of ViT on the WCE dataset. Excitingly, our method successfully achieves the best performance, and this proves the usability of ViT on the WCE dataset. Further work on classification and diagnosis on the WCE dataset will aim to reduce the number of parameters of ViT and speed up the inference speed, so that ViT can be better implemented in the real-world scenario of clinical medical treatment.

6. Conclusions

This paper presents a ViT-based WCE image disease classification algorithm that can be trained from scratch. We adopt several CNN and ViT algorithms for comparison experiments. The experimental results show that, in terms of accuracy, our adopted methodology surpasses all comparison methods in processing small-scale imbalanced WCE datasets, demonstrating the feasibility of applying the ViT algorithm directly to small-scale medical datasets. Our proposed method can derive classification results with the inference time of 1 ms, making it suitable for real-time diagnosis. However, although our model achieves greater accuracy than existing methods on two public datasets, the slightest misdiagnosis or miss in medical applications can be fatal. Thus, our solution can only be used for prescreening before the doctor’s manual diagnosis, or for review after the doctor’s diagnosis. AI still needs further exploration in actual medical diagnosis and treatment.