A Two-Step Learning Model for the Diagnosis of Coronavirus Disease-19 Based on Chest X-ray Images with 3D Rotational Augmentation

Abstract

Herein, we propose a method for effectively classifying normal, coronavirus disease-19 (COVID-19), lung opacity, and viral pneumonia symptoms using chest X-ray images. The proposed method comprises a lung detection model, three-dimensional (3D) rotational augmentation, and a two-step learning model. The lung detection model is used to detect the position of the lungs in X-ray images. The lung position detected by the lung detection model is used as the bounding box coordinates of the two-step learning model. The 3D rotational augmentation, which is a data augmentation method based on 3D photo inpainting, solves the imbalance in the amount of data for each class. The two-step learning model is proposed to improve the model performance by first separating the normal cases, which constitute the most data in the X-ray images, from other disease cases. The two-step learning model comprises a two-class model for classifying normal and disease images, as well as a three-class model for classifying COVID-19, lung opacity, and viral pneumonia among the diseases. The proposed method is quantitatively compared with the existing algorithm, and results show that the proposed method is superior to the existing method.

1. Introduction

The coronavirus disease-19 (COVID-19) pandemic has spread worldwide and significantly affected almost all countries and regions. According to statistics by the World Health Organization, in March 2022, more than 480 million infections and more than 6 million deaths were reported [1]. COVID-19 is an infectious disease caused by the SARS-CoV-2 virus. Most patients with COVID-19 exhibit mild-to-severe symptoms and recover from the disease without specific treatment. However, in some cases, the disease progresses to a serious condition that necessitates medical attention. When an infected person coughs, sneezes, or speaks and comes into contact with another person within 2 m, droplets of phlegm are produced that can cause infection when they are transmitted. In general, the average incubation period for an infected person is 5–6 days, and total recovery may require 14 days. The main symptoms of COVID-19 include a fever of 37.5 °C or higher, cough, fatigue, and loss of smell and taste sensations [2]. Prompt and accurate testing must be performed to prevent the spread of the disease. Currently, reverse transcription polymerase chain reaction (RT-PCR) [3] is used for the diagnosis of COVID-19; however, RT-PCR involves a long test time. Chest X-ray (CXR) and computed tomography are typical techniques for diagnosing lung-related diseases such as pneumonia and tuberculosis, and they can also be used to test for COVID-19 [4,5]. CXR images can be used to identify diseases occurring in the lungs more rapidly and at a lower cost compared with other methods.

Recently, studies based on CXR imaging for detecting COVID-19 have been actively conducted. Deep learning models, which are excellent for object detection and classification, have been used extensively. Ioannis et al. [6] used 1427 X-ray datasets to classify COVID-19 and normal X-ray images, which afforded an accuracy of 96.78% in MobileNetv2 [7] with transfer learning. Abbas et al. [8] proposed decomposing, transferring, and composing (DeTraC) to classify COVID-19 and normal X-ray images using a dataset comprising 196 X-ray images. DeTraC was designed to investigate class boundaries using a class decomposition mechanism to manage irregularities occurring in the dataset and it afforded an accuracy of 95.12%. Minaee et al. [9] classified COVID-19 and non-COVID-19 using a dataset containing only 184 COVID-19 X-ray images from a total of 5184 images. They applied transfer learning and fine tuning to existing convolutional neural network (CNN) models, i.e., residual networks (ResNet)-18, ResNet50 [10], SqueezeNet [11], and densely connected convolutional networks (DenseNet) [12], where an accuracy of 98% (±3%) was indicated for each model. Sensitivity and specificity levels of approximately 90% were achieved. Panwar et al. [13] proposed the nCOVnet that comprises five custom layers based on the visual geometry group (VGG)-16 model. The nCOVnet used a dataset comprising 142 COVID-19 X-ray images and 142 normal images captured with posterior anterior views; it achieved sensitivity and specificity levels of 97.62% and 78.57%, respectively. Keidar et al. [14] used 2462 CXR images and a pre-trained ResNet50 model as well as applied data augmentation and lung segmentation as preprocessing algorithms to improve model performance. A new three-channel image comprising original CXR, lung segmentation image, and zero image was used as the training data for the model. The accuracy and sensitivity were 89.7% and 87.1%, respectively. Wang et al. [15] used 13,975 CXR datasets obtained from five open-access data repositories for training and proposed COVID-Net, which comprised a lightweight residual projection–expansion–projection–extension design pattern. COVID-Net classifies three classes, i.e., normal, non-COVID-19, and COVID-19, and affords an accuracy of 93.3%. Afshar et al. [16] proposed the COVID-CAPS that comprises four convolutional layers and three capsule layers; it afforded accuracy, sensitivity, and specificity of 95.7%, 90%, and 95.8%, respectively. Khan et al. [17] proposed a model for classifying normal lung opacity, COVID-19, and viral pneumonia using 21,165 CXR images. The authors applied data augmentation such as rotation and flip to solve imbalance in the amount of data for each class. The head section of the existing pre-trained CNN model was redesigned with batch normalization, a dropout layer, and two dense layers to achieve an accuracy of 96.13%.

Contrary to the use of original CXR images in existing methods, images obtained by segmenting only the lung region among CXR images have been used in a few studies [18,19]. When the results of a CNN model are interpreted using visualization methods such as Score-CAM [20], most of the regions contributing to the decision of the CNN model are not the lung region. Although the segmentation of the lung region may not improve the model performance, the reliability of the model results must be improved. Rahman et al. [19] compared the performance of a CNN model based on various data augmentation and lung segmentation methods using the COVQU-20 dataset that comprises 18,479 normal, non-COVID-19 lung opacity, and COVID-19 CXR images. Data augmentation uses gamma correction, complement, histogram equalization, contrast-limited adaptive histogram equalization, and balanced contrast enhancement techniques that vary the brightness and contrast of images. Meanwhile, Resnet18, Resnet50, Resnet101, CheXNet [21], and DenseNet201 were used as CNN models. U-Net [22] was used to segment the lung region in the image. In the plain CXR dataset, gamma correction and CheXNet performed the best, whereas in the lung-segmented dataset, gamma correction and DenseNet201 performed the best. The accuracy, precision, and recall achieved were 96.29%, 96.28%, and 96.28%, respectively, in the non-segmented case, and 95.11%, 94.55%, and 94.56%, respectively, in the segmented case. Aslan at el. [23,24] used 2905 CXR images to classify normal, COVID-19, and viral pneumonia. Artificial neural networks (ANN)-based segmentation model is used to crop the lung region from raw CXR, obtain robust features, and improve classification accuracy. The ANN model consists of a single hidden layers with 100 neurons and output with 6 neurons. The output consists of x and y coordinates of three points that are converted to lung region. In [23], they proposed CNN-based transfer learning bidirectional long short-term memories model with an accuracy 98.70%. In article [24], they compared five pre-trained CNN models to extract the features and four machine learning (ML) algorithms to classify the CXR images. To improve the accuracy, Bayesian optimization was applied the ML algorithms. They achieved the highest accuracy 96.29% with DenseNet201 model and support vector machine algorithm. The sensitivity, precision, specificity, F1 score, and Matthews correlation coefficient (MCC) of this model are 0.9642, 0.9642, 09812, 0.9641, and 0.9453, respectively.

Herein, we propose a deep learning model that can effectively classify normal, COVID-19, lung opacity, and viral pneumonia from a COVID-19 dataset and reduce false-negative COVID-19 diagnoses to prevent the spread of infection in the surroundings. Our goal is to propose a new data augmentation method to overcome the problem due to imbalanced datasets, and a two-step learning model for effective separation of normal and disease CXR images and the classification of diseases from each disease CXR image. The proposed method involves a lung detection model, 3D rotational augmentation, and a two-step learning model. The lung detection and two-step learning models are designed based on you only look once (YOLO)-v4 algorithm [25]. The training image and coordinates of the bounding box corresponding to each image are required to train YOLOv4. However, denoting the coordinates of the bounding box manually in numerous images is time consuming. Hence, a lung detection model is proposed to detect the position of the lung and automatically generate the coordinates of the position of the lung as a label file. Additionally, 3D rotational augmentation is proposed to balance data by increasing the number of non-uniform data for each class based on 3D photo inpainting (3DPI). Unlike the existing two-dimensional (2D) image rotation, 3D rotational augmentation shifts the viewpoint of an object photographed with a camera such that an image with a more realistic shape can be created. The two-step learning model comprises a two-class model that classifies normal and disease images and a three-class model that classifies COVID-19, lung opacity, and viral pneumonia images among the diseases classified in the two-class model. A two-step learning model is proposed to prevent erroneous detection of the normal image that constitutes the largest proportion of the COVID-19 dataset in other classes. A quantitative comparison of the proposed method with the existing algorithm indicates the superiority of the proposed method over existing methods.

2. Related Studies

2.1. YOLOv4

YOLO [26] is a representative object-detection algorithm. YOLO predicts multiple bounding box coordinates and class probabilities simultaneously in a single convolutional network without using region proposals, which are used to locate objects in existing R-CNN models [27,28,29]. Therefore, the processing speed of YOLO is much higher than that of existing algorithms and real-time object detection is realizable. YOLOv2 [30] improves the model performance using batch normalization, anchor box, direct location prediction, and multiscale training methods. YOLOv3 [31] uses a bounding box with a multiscale and residual block to solve the problem of small object detection in the existing YOLO. YOLOv4 improves the model performance by applying cross-stage-partial-connection [32], spatial pyramid pooling (SPP) [33], and path aggregation network (PAN) [34]. YOLOv4 comprises a backbone, neck, and head. The backbone of YOLOv4 is CSPDarknet53 that comprises batch normalization and Mish activation functions. In CSPDarknet53, only one-half of the feature map passes through the residual block; therefore, a bottleneck layer is not required. The neck connects the backbone to the head and reconstructs the feature map. The neck of YOLOv4 is composed of SPP and PAN. SPP effectively increases the receptive field using contextual features extracted from four max-pooling layers. The PAN improves the overall feature layers and localization performance by expanding the path used for information flow between the lower layer and uppermost feature. The head has the same structure as YOLOv3 and predicts the prediction boxes using feature maps of three different scales. In addition, data augmentation using a new technique known as mosaic is used in YOLOv4 to improve the object detection ability of the training data. Figure 1 shows the architecture of YOLOv4.

2.2. 3DPI

3D image conversion refers to rendering a 2D image captured by a camera into an image of a new view. It can be used to record and reproduce visual perceptions from various angles. Classical image-based rendering technology requires sophisticated image capturing technology using specific hardware. Recently, 3D image conversion methods using RGB-D (red, green, and blue-depth) images obtained from a small digital camera or mobile phones equipped with a dual lens instead of expensive customized equipment have been investigated.

In 3DPI, parallax is separated using an RGB-D image and the information lost when an image of a new rendered view is restored. The 3DPI technique can be classified into image-based and learning-based rendering. In image-based rendering, an image of a new view is synthesized using a set of posed images. It achieves good performance when the multiview stereo algorithm functions well or the image is captured using a depth sensor. Recently, many learning-based rendering techniques based on CNNs have been investigated.

The learning-based rendering technique does not require expensive equipment because it can synthesize a new view image using a single image and a stereo pair image. Shih et al. [35] used a layered depth image (LDI) to reduce the complexity of arbitrary depth information. The input LDI image is classified into several local regions while considering the connectivity between pixels, and the image synthesized by applying the inpainting algorithm is merged into one LDI image. Subsequently, the depth information is changed and applied repeatedly. Shih’s algorithm comprises three subnetworks: color inpainting, depth inpainting, and edge inpainting. The edge inpainting network preferentially restores the edge of the region requiring restoration. After edge restoration is performed, the color and depth information of the image boundary region is restored from the color and depth inpainting network.

2.3. Data Augmentation

Deep-learning models require significant amounts of data to achieve effective learning. However, the data required for learning are difficult to obtain and the obtained data may cause an imbalance in the amount of data for each class. When a deep-learning model is trained with such data, it can be biased toward a specific class, thereby significantly affecting the model performance. Data augmentation is performed to overcome this problem. Because data augmentation can increase the number of training data in various methods, it can address any insufficient data or imbalance for each class. Data augmentation techniques include geometric transformation, color space transformation, kernel filter-based augmentation, random erasing, image mixing, and mosaic [36]. Geometric transformation is a method of changing the shape of an image by rotating, translating, scaling, cropping, and flipping the image. Color space transformation is a method of changing the color or brightness of an RGB (red, green, and blue) image by adjusting the pixel values. It is used to compensate for biased data due to limited lighting environments. Kernel-filter-based augmentation can generate blur or sharpen images by adjusting the internal value of an n × n filter. Random erasing refers to a method of inserting a patch with a random size, position, and value into an image [37]. In this method, dropout regularization is implemented in the input data space instead of in the network to improve the performance degradation of the model caused by occlusion between objects in the image. Meanwhile, mixing images is a method of synthesizing images of different classes with the average pixel value and it is performed in the classification model [38]. Mosaic refers to combining four images used for training into one image at a specific ratio [25]. Mosaics can learn to identify objects on a scale smaller than the normal scale, thereby improving the recognition ability of models in complex backgrounds. In addition, it can significantly reduce the necessity for a large mini-batch size during training.

3. COVID-19 Dataset

The COVID-19 dataset included images published by Kaggle [19,39,40]. The data points were categorized into four classes: normal, COVID-19, lung opacity, and viral pneumonia. Each class was created by merging several previously published data sets. The normal data points comprised 8851 images from the Radiological Society of North America (RSNA) [41] and 1341 images from Kaggle’s chest X-ray images (pneumonia) [42]. The total number of images for COVID-19 was 3616, and each image was from the BIMCV-COVID19 dataset, Germany Medical School, Italian Society of Medical Radiology (SIRM), Github, Kaggle, Twitter, and the COVID-19 CXR repository [43,44,45,46,47,48,49,50,51]. The lung opacity data comprised 6012 CXR images acquired from the RSNA. The number of viral pneumonia data points was 1345 and they were obtained from CXR images (pneumonia). The CXR images of the dataset were in grayscale. The resolution of each image was 299 × 299 pixels and the image format was portable network graphics. Figure 2 shows the sample images for each class.

Table 1. Composition of COVID-19 dataset.

Class	RSNA	CXR Images (Pneumonia)	BIMCV-COVID19+	Germany Medical School	SIRM, Github, Kaggle, and Twitter	COVID-19 CXR Repository	Total
Normal	8851	1341					10,192
COVID-19			2473	183	560	400	3616
Lung Opacity	6012						6012
Viral Pneumonia		1345					1345

presents the number of images in each class.

4. Proposed Method

Herein, we propose an algorithm to effectively classify normal, COVID-19, lung opacity, and viral pneumonia based on CXR images. The proposed deep learning model is based on YOLOv4. A block diagram of the proposed algorithm is shown in Figure 3. The proposed model comprises a lung detection model for labeling automation, dataset preparation for data augmentation and labeling, and a two-step learning model for CXR classification. The two-step learning model is used to improve the detection performance of CXR. The two-step learning model comprises a two-class model for classifying normal and disease images, and a three-class model to classify COVID-19, lung opacity, and viral pneumonia among the diseases. Additionally, the two- and three-class models are used to detect normal cases and each disease separately. The input CXR is classified as normal or disease in the two-class model. The disease cases detected by the two-class model are classified as COVID-19, lung opacity, or viral pneumonia in the three-class model.

4.1. Lung Detection Model

The input image and label information corresponding to each image are required to train the object detection algorithm. The location and class information of the object are included in the label file and expressed numerically. The object location information is expressed in terms of the (x, y) coordinates, as well as the width and height in a bounding box. Here, (x, y) are the center coordinates of the bounding box. Each value ranges from 0 to 1. The label file used for training is created manually in general; therefore, considerable time and effort is required. A lung detection model is proposed for increasing the efficiency of the operation by automatically generating a label file. The proposed model aims to detect the position of the lungs using minimal training data from the entire image. The processing of the proposed lung detection model is shown in Figure 4.

The proposed lung detection model is trained twice. The first learning method uses only 25 images from each class of the COVID-19 dataset. After the training is completed, the first model is used to detect the position of the lungs. If two lungs are not detected, then the image is registered as an error image. The second training is performed using 100 registered error images, as well as the 100 images used in the first model.

4.2. 3D Rotational Augmentation

The COVID-19 dataset is disproportionate, with the ratio of normal, COVID-19, lung opacity, and viral pneumonia at approximately 8:3:4:1. To improve the model performance, the amount of data must be balanced by increasing the number of images for each class.

Therefore, we propose a 3D rotational augmentation method for data augmentation based on 3DPI [35]. In 3D rotation augmentation, a virtual depth image is applied to a 2D plane, resulting in rotating an image in a 3D direction. When capturing X-ray images, the possibility of imaging while the patient’s posture is turned left or right should be considered.

In 3DPI, a depth map is required to determine the effect of 3D photography by shifting an object. The depth map separates the foreground and background from the images based on depth information. A depth map is generally created using two stereo pair images. However, because only one image exists in the COVID-19 dataset, the depth-map generation algorithm cannot be easily applied using stereo vision. Therefore, a simple depth map applicable to 3DPI is proposed. The proposed depth map is created by extracting the brightness of the horizontal line at the center of the chest from the average of multiple CXR images and then copying the same value in the vertical direction. The depth map is the brightest in the center of the chest and the brightness gradually decreases toward the lungs.

Figure 5a shows the results of the proposed depth map and MiDas [52], which is a depth-estimation network used in 3DPI. Figure 5b shows the results of applying the 3DPI algorithm to each depth map. Although the depth map generated using MiDas shows the form of a body, a distortion was observed in some areas (red and blue boxes), unlike the original image. However, the proposed method does not result in image distortion and the body position shifts naturally. Figure 6 shows the resulting image of 3D rotational augmentation. Figure 6a shows the direction of image rotation. The rotation directions are counterclockwise and clockwise with respect to the center of the image. Figure 6b shows the images generated based on each rotational direction. The center of the image is denoted by a dotted red line and the position of the esophagus is denoted by dotted black line to display the changed image. The black arrow in each image indicates the distance shifted from the image center. The results of the 3D rotational augmentation show that the size of the left and right lungs varies depending on the rotation degree, unlike those of shift and horizontal flipping.

4.3. Two-Step Learning Model

The two-step learning model comprises a two-class model for classifying normal and disease images and a three-class model for classifying COVID-19, lung opacity, and viral pneumonia among the diseases. Additionally, 3D rotational augmentation is used to match the number of data points in each class uniformly before training. A label file is created using the lung detection model. The position and class information of the left and right lungs are included in the label file. In the two-class model, the normal class ID is set to 0 and the IDs of all other classes are set to 1. In the three-class model, the class IDs are set to 0, 1, and 2 for COVID-19, lung opacity, and viral pneumonia, respectively. The two- and three-class models are trained using YOLOv4. A block diagram of the proposed two-step learning model is shown in Figure 7.

In the detection step, the two- and three-class models are used sequentially. First, normal and disease images are classified using the two-class model, and among the diseases classified in the two-class model, COVID-19, lung opacity, and viral pneumonia are classified in the three-class model. Figure 8 shows a block diagram of the detection process, and Figure 9 shows the resulting images yielded by the proposed two- and three-class models. In the resulting images, both the left and right lung positions are accurately detected; meanwhile, normal and disease images are detected accurately in the two-class model, in addition to the class of each disease in the three-class model.

Multiple objects are detected during the detection stage of the two-step learning model. The maximum class sum (MCS) method was used to classify one symptom from the multiple detection results. MCS is a method of selecting the class with the largest sum after calculating the confidence sum for each detected class. The maximum class sum is expressed as follows:

$$\mathrm{MCS}=\underset{i}{\max }{\displaystyle \sum }_{j=1}^n{C}_{ij},$$

(1)

where $\mathrm{MCS}$ represents the class selected based on the sum of the maximum classes; $i$ and $j$ are the classes of the detected bounding box and the number of detected bounding boxes, respectively; and $C$ is the confidence level of the bounding box. If the detected bounding box did not exist, it was denoted as “N/A.”

5. Simulation and Results

The system used in the experiment was configured with an NVIDIA 3090 GPU, an i9-10980EX CPU, 256 GB of RAM, CUDA 11.0, and cuDNN 8.0.4. YOLOv4 uses a C-language-based Darknet framework [53]; meanwhile, 3DPI uses PyTorch 1.9.0 and Python 3.8.10. The size of the input image used for training YOLOv4 was 608 × 608 pixels; batch size was 64, learning rate was 0.001, momentum was 0.949, and the decay was 0.005. CSPDarkNet53 was used as the backbone.

Six performance metrics, i.e., overall accuracy, weighted precision, weighted recall, weighted F1 score, weighted area under curve (AUC), and weighted MCC were used to evaluate the model performance. Each performance metric is expressed as follows:

$$\mathrm{Overall}\mathrm{Accuracy}\left(\mathrm{OA}\right)=\frac{TP+TN}{TP+FN+FP+TN}$$

(2)

$$\mathrm{Weighted}\mathrm{Precsion}\left(\mathrm{WP}\right)=\frac{{\sum }_{c\in C}{N}_c\left(\frac{T{P}_c }{T{P}_c+F{N}_c}\right)}{{\sum }_{c\in C}{N}_c}$$

(3)

$$\mathrm{Weighted}\mathrm{Recall}\left(\mathrm{WR}\right)=\frac{{\sum }_{c\in C}{N}_c\left(\frac{T{P}_c}{T{P}_c+F{N}_c}\right)}{{\sum }_{c\in C}{N}_c}$$

(4)

$$\mathrm{Weighted}\mathrm{F}1\mathrm{Score}\left(\mathrm{WF}1\right)=\frac{{\sum }_{c\in C}{N}_c\left(\frac{2\times T{P}_c}{2\times T{P}_c+F{N}_c+F{P}_c}\right)}{{\sum }_{c\in C}{N}_c}$$

(5)

$$\mathrm{Weighted}\mathrm{AUC}\left(\mathrm{WAUC}\right)=\frac{{\sum }_{c\in C}{N}_c\left(\frac{TP}{TP+FN}+\frac{TN}{TN+FP}\right)}{2{\sum }_{c\in C}{N}_c}$$

(6)

$$\mathrm{Weighted}\mathrm{MCC}\left(\mathrm{WMCC}\right)=\frac{{\sum }_{c\in C}{N}_c\left(\frac{\left(TP\times TN\right)-\left(FN\times FP\right)}{\sqrt{\left(TP+FN\right)\times \left(TN+FP\right)\times \left(TP+FP\right)\times \left(TN+FN\right)}}\right)}{{\sum }_{c\in C}{N}_c}$$

(7)

In the above, $TP$, $TN$, $FP$, and $FN$ represent true positive, true negative, false positive, and false negative, respectively; $C$ is the class of the CXR dataset and $N_c$ is the number of test CXR images in class $C$.

5.1. Lung Detection Model

The lung-detection model was trained twice. In terms of the training data of the first model, 100 images were used, with 25 images randomly selected for each class of normal, COVID-19, lung opacity, and viral pneumonia.

The label information consisted of the position of the left and right lungs in each image. The ratio of training and validation datasets was 80:20, and only one class, i.e., “lung”, was available. The internal default settings of YOLOv4 were used for data augmentation. Because all CXR images were grayscale images, hue and saturation augmentation were not performed. After learning was completed, all CXR images were evaluated, and the images that did not detect the position of the two lungs were classified as error images. Retraining was performed using 100 cases of the classified error images in addition to the 100 images used in the first training. The training conditions were set to be the same as those in the first training.

Table 2. Comparison results between first and second lung detection models.

Lung Detection Model	Number of Bounding Boxes					Errors	Error Rate
	Ground Truth	Detection Results
	2	0	1	2	3
Frist model (Using 100 training images)	21,165	8	158	21,007	103	158	0.747%
Second model (Using 200 training images)	21,165	0	35	21,130	32	35	0.165%

shows the number of bounding boxes in the first and second lung detection models for CXR images. In the first model, 158 false detections occurred; however, in the second model, 35 false detections were indicated.

5.2. COVID-19 Classification

The results of the proposed model were compared with those of VGG16 [54] and YOLOv4. VGG16 comprises 13 convolutional layers and 3 fully connected layers. To train VGG16 using the COVID-19 dataset, the output of the last fully connected layer was modified to four and retrained using pre-trained weights with the ImageNet dataset. Flip, exposure, shift, and scale augmentation were applied to VGG16. YOLOv4 with default augmentation, except for hue and saturation, and the maximum class sum, were applied to classify the results. In both the models, additional data augmentation was not applied; the datasets used for training and testing are listed in

Table 3. Details of imbalanced dataset used for training, validation, and test.

Class	Training Image	Validation Image	Test Image
Normal	8392	200	1600
COVID-19	1816	200	1600
Lung Opacity	4212	200	1600
Viral Pneumonia	945	200	200

.

The results of the proposed two-step learning model were classified into two datasets: imbalanced and balanced. First, the two-step learning model was trained using the original imbalanced dataset shown in Table 3 without external augmentation. Subsequently, the two-step learning model was trained using the balanced dataset via 3D rotational augmentation. In a balanced dataset, the number of training images in each class is matched. As shown in

Table 4. Details of balanced dataset used for training, validation, and test.

Class	Training Image	Validation Image	Test Image
Normal	24,000	6000	1600
COVID-19	8000	2000	1600
Lung Opacity	8000	2000	1600
Viral Pneumonia	8000	2000	200

, the ratio of normal to disease was 1:1, and the ratio of COVID-19, lung opacity, and viral pneumonia was 1:1:1. The imbalanced two-step learning model uses the default augmentation of YOLOV4, and the balanced two-step learning model applies 3D rotational augmentation instead of flipping and exposure augmentation.

Figure 10 shows the charts of training loss and mean average precision (mAP) for two-and three-class models. The average loss and mAP for the two-class model is 0.5886 and 100.0% and those of three-class model are 0.6586 and 99.5%.

Table 5. Comparison results for four learning cases.

Model	Performance Metric
	OA	WP	WR	WF1	WAUC	WMCC
VGG16 (Imbalanced)	83.42	84.69	83.42	83.41	87.51	75.83
YOLOv4 (Imbalanced)	88.26	90.03	88.26	88.44	91.31	83.58
Two-step learning model (Imbalanced)	91.92	92.26	91.92	91.94	94.04	88.26
Two-step learning model with 3D rotational augmentation (Balanced)	95.40	95.72	95.40	95.49	96.71	93.46

and Figure 11 show the performance of each model. When classifying four simple classes, the accuracy of VGG16 was 83.42%, and that of YOLOv4 was 88.26%, thus confirming the excellent performance of YOLOv4. The accuracy of the imbalanced two-step learning model for the same dataset was 91.26%, which was higher than that of YOLOv4. The overall accuracy, weighted precision, weighted recall, weighted F1 score, weighted AUC, and weighted MCC of the balanced two-step learning model using 3D rotational augmentation were the highest at 95.4, 95.72, 95.40, 95.49, 96.71, and 93.46%, respectively. Therefore, the proposed two-step learning model with 3D rotational augmentation improved the detection performance.

Table 6. Comparison results of detection for normal, COVID-19, lung opacity, and viral pneumonia.

Model	Accuracy
	Normal	COVID-19	Lung Opacity	Viral Pneumonia
VGG16 (Imbalanced)	94.13	75.69	81.06	78.50
YOLOv4 (Imbalanced)	98.81	84.44	82.31	82.00
Two-step learning model (Imbalanced)	93.44	96.38	85.63	94.50
Two-step learning model with 3D rotational augmentation (Balanced)	97.38	96.50	91.75	100.00

and Figure 12 show the comparison results of detection for normal, COVID-19, lung opacity, and viral pneumonia. The disease determination accuracy for COVID-19 patients afforded by the proposed two-stage learning model exceeded 96%. When balanced data were used, the judgment error of the normal case was less than 3%.

6. Discussion

This study aimed to effectively classify normal, COVID-19, lung opacity, and viral pneumonia using CXR imaging. The COVID-19 CXR dataset has an imbalanced number of images for each disease, and the number of normal CXR images is greater than the number of disease CXR images. In order to improve the performance of the model, it is necessary to match each dataset ratio using data augmentation. The existing data augmentation such as image rotation does not reflect the shape of the actually photographed CXR images. In general, because CXR images are taken around the center of the body, 3D rotational augmentation is more effective than 2D image rotation. Therefore, the 3D rotational augmentation method using 3DPI was proposed. Then, the two-step learning model was proposed to improve the model performance by first separating the normal cases, which constitute the highest amount of data in the X-ray images, from other disease cases. Figure 13 shows the confusion matrix of each method. Figure 13a,b shows the results of VGG16 (Imbalanced) and YOLOv4 (Imbalanced). The accuracy of VGG16 (Imbalanced) for normal, COVID-19, lung opacity, and viral pneumonia was 94.13, 75.69, 81.06, and 78.05%, and YOLOv4 was 98.81, 84.44, 82.31, and 82.00%. This result indicates that the performance of YOLOv4 is better than that of VGG16. In Figure 13c, the normal accuracy of the two-step learning model (Imbalanced) is 93.44%, which is lower than that of YOLOv4, but the accuracy of COVID-19, lung opacity, and viral pneumonia is 96.38, 85.63, 94.50%, which is higher than that of YOLOv4. This confirms that the three-class model classifies diseases well. The lowered accuracy of normal is predicted to occur due to data imbalance during training of the two-class model. Figure 13d represents the two-step learning model with 3D rotational augmentation (Balanced). The accuracy of the model for normal, COVID-19, lung opacity, and viral pneumonia is 97.38, 96.50, 91.75, and 100.00%, which is the best. This confirms that the proposed 3D rotational augmentation effectively improved the performance of the model. However, the two-step learning model has a problem that requires a lot of training time because two models need to be trained. In the final result compared to the imbalanced YOLOv4 model, additional research is needed on the fact that the accuracy of normal is lowered from 98.81% to 97.38%.

The proposed method has a higher WR score than those of existing methods. This indicates that the proposed model has fewer FN cases. An FN case is a failure to identify an infected patient. A highly contagious disease such as COVID-19 can easily spread through a community due to misdiagnosis. Therefore, there is a need for a method to reduce FN cases. The proposed method will be helpful in clinical decisions because the proposed method has better performance than the existing method.

7. Conclusions

Herein, we proposed a deep learning model that could effectively classify normal, COVID-19, lung opacity, and viral pneumonia images in an imbalanced COVID-19 dataset. The proposed method comprises a lung detection model for labeling automation, 3D rotational augmentation to solve the imbalance of insufficient data, and a two-step learning model to enhance the model performance. The lung detection model could accurately detect the location of the lungs in the CXR using only 200 images from a dataset comprising approximately 20,000 images. The 3D rotation augmentation created a more realistic image by shifting the viewpoint of the object captured by the camera. The two-step learning model performed better than the existing method of merely classifying the classes into one of four categories by first separating the normal images that constituted the largest proportion of imbalanced data from the other diseases. The two-step learning model with 3D rotational augmentation performed better than the existing method in terms of overall accuracy, weighted precision, weighted recall, weighted F1 score, weighted AUC, and weighted MCC with values of 95.40, 95.72, 95.40, 95.49, 96.71, and 93.46%, respectively. The two-step learning model afforded an accuracy exceeding 96% for COVID-19 patients. The proposed method is expected to serve as a rapid screening tool during the COVID-19 pandemic.