PneuNet: deep learning for COVID-19 pneumonia diagnosis on chest X-ray image analysis using Vision Transformer

Abstract

A long-standing challenge in pneumonia diagnosis is recognizing the pathological lung texture, especially the ground-glass appearance pathological texture. One main difficulty lies in precisely extracting and recognizing the pathological features. The patients, especially those with mild symptoms, show very little difference in lung texture, neither conventional computer vision methods nor convolutional neural networks perform well on pneumonia diagnosis based on chest X-ray (CXR) images. In the meanwhile, the Coronavirus Disease 2019 (COVID-19) pandemic continues wreaking havoc around the world, where quick and accurate diagnosis backed by CXR images is in high demand. Rather than simply recognizing the patterns, extracting feature maps from the original CXR image is what we need in the classification process. Thus, we propose a Vision Transformer (VIT)-based model called PneuNet to make an accurate diagnosis backed by channel-based attention through X-ray images of the lung, where multi-head attention is applied on channel patches rather than feature patches. The techniques presented in this paper are oriented toward the medical application of deep neural networks and VIT. Extensive experiment results show that our method can reach 94.96% accuracy in the three-categories classification problem on the test set, which outperforms previous deep learning models.

Keywords: COVID-19; Deep learning; Multi-head attention; Pneumonia diagnosis; Vision Transformer.

Fig. 1

Architecture of PneuNet (a) and the details of Transformer Encoder (b)

Fig. 2

Architecture of ResNet18, without flatten layer nor logit layer

Fig. 3

Image partition in conventional VIT process: raw image (a) is divided into several patches (b) and each patch will be encoded and embedded with its position

Fig. 4

Illustration of patches employed in Transformer modular which obtained from ResNet18

Fig. 5

Typical chest X-ray images from the combined dataset: CXR image of None Pneumonia (a), CXR images of COVID-19 (b), CXR images of Bacterial Pneumonia (c), and CXR images of Viral Pneumonia (d)

Fig. 6

History during training: history of cross-entropy loss (a) and history of categorical accuracy (b)

Fig. 7

The confusion matrix for three-category classification: None Pneumonia, Normal Pneumonia, and COVID-19

Fig. 8

Confusion matrix generated from binary classification model (a) and four-category classification model (b)