GIONet: Global information optimized network for multi-center COVID-19 diagnosis via COVID-GAN and domain adversarial strategy

Abstract

The outbreak of coronavirus disease (COVID-19) in 2019 has highlighted the need for automatic diagnosis of the disease, which can develop rapidly into a severe condition. Nevertheless, distinguishing between COVID-19 pneumonia and community-acquired pneumonia (CAP) through computed tomography scans can be challenging due to their similar characteristics. The existing methods often perform poorly in the 3-class classification task of healthy, CAP, and COVID-19 pneumonia, and they have poor ability to handle the heterogeneity of multi-centers data. To address these challenges, we design a COVID-19 classification model using global information optimized network (GIONet) and cross-centers domain adversarial learning strategy. Our approach includes proposing a 3D convolutional neural network with graph enhanced aggregation unit and multi-scale self-attention fusion unit to improve the global feature extraction capability. We also verified that domain adversarial training can effectively reduce feature distance between different centers to address the heterogeneity of multi-center data, and used specialized generative adversarial networks to balance data distribution and improve diagnostic performance. Our experiments demonstrate satisfying diagnosis results, with a mixed dataset accuracy of 99.17% and cross-centers task accuracies of 86.73% and 89.61%.

1. Introduction

Based on a study [1], the pneumonia caused by the novel coronavirus (COVID-19) is a severe acute illness, and the original strain in early 2020 had a mortality rate of nearly 2.0%. Once a person contracts COVID-19, they usually develop respiratory symptoms within a few days, which can quickly progress to acute respiratory distress syndrome. The rapid progression of the disease makes it difficult to detect in patients with community-acquired pneumonia.(CAP). According to the latest report from the World Health Organization [2], as of now, COVID-19 has resulted in 624 million confirmed cases and 6.5 million deaths globally, imposing significant health risks and economic burdens worldwide. Although the current mainstream strain, Omicron, has significantly reduced mortality rates, its contagiousness is astonishing and can still cause a huge number of new COVID-19 cases. Early screening of COVID-19 cases is critical for effective prevention and treatment of the disease. Detecting COVID-19 patients early and providing prompt treatment can effectively prevent disease progression and significantly reduce mortality rates. Therefore, finding accurate and efficient screening methods for COVID-19 patients is essential. Reverse transcription polymerase chain reaction (RT-PCR) has been widely used for COVID-19 screening, however, it is still limited by the large consumption of medical resources [3].

The typical medical diagnostic methods for pulmonary diseases include X-rays and computed tomography scans (CT) [4], so it can be inferred that these two medical imaging technologies are also feasible auxiliary diagnostic scheme for COVID-19 [5]. Although X-ray imaging can identify bronchitis, bronchiolitis, and shadows in the lung field, its rate of missed diagnosis is high, and there are no significant image features in the early stage of the disease [6]. On the other hand, ground-glass opacities (GGO) in the CT image are a typical manifestation that can confirm COVID-19 [7]. GGO and infiltrative shadows appear in both lungs with the development of the disease, and pulmonary consolidation may occur in severe cases [8]. By analyzing the distribution and size of lesions in CT images, a pulmonary inflammation index can be mapped to assess the severity of lung inflammation [9]. However, manual interpretation requires a significant amount of experienced medical personnel and is time-consuming, which is not feasible during the COVID-19 pandemic when medical resources are severely limited [10]. Therefore, it is crucial to use deep learning methods for the automatic diagnosis of COVID-19 [11].

Currently, many researchers are using deep learning to assist in the study of COVID-19. For example, Ardakani et al. [12] utilized ten typical Convolutional Neural Networks (CNNs) for the joint diagnosis of COVID-19 on transverse CT images. The above methods improve the diagnostic performance of the model by cascading multiple neural networks, but training is cumbersome, the number of parameters is huge, and the final result is only obtained by integrating the results of multiple models. Similarly, Ko et al. [13] has deployed a fast-track COVID-19 classification network for conducting a three-class study on COVID-19 using transfer learning from four advanced networks. In addition to the aforementioned studies, Shah et al. [14] and Wang et al. [15] also conducted methodological research on two-dimensional convolutional neural networks. However, 2D slice images may lose the relevance of axial features, and multiple slices of the same case at different locations may conflict during the classification process. How to automate the selection of slices with diagnostic significance is a challenge. Therefore, it can be seen that conducting research on 3D volume images is more clinically significant. For example, Serte et al. [16] used 3D CNN to diagnose COVID-19 and achieved a 96.0% AUC score. In this study, only a single backbone network was used and no comparison was made with different networks. In addition, He et al. [17] utilized the differentiable neural architecture search method to automatically search for the optimal network structure for the diagnosis of COVID-19. Which has the advantages of automation, efficiency and flexibility, but suffers from high computational cost, data requirements, and the possibility of suboptimal results due to the complexity of the search process.

Due to COVID-19 being a global pandemic, multi-center research and even cross-national research is urgently needed. However, there are currently only limited studies that focus on multi-center research. One study [18] utilized clinical data and CT images to predict the mortality rate of patients across four centers. However, they only utilized data from different centers for classification and did not address the feature differences present in multi-center data. Another study [19] introduced a new robust weakly supervised learning approach to address the variation in CT image features among different centers. Due to the heterogeneity of CT images from different centers as described in the above-mentioned studies, it is necessary to address this heterogeneity. Therefore, we plan to use domain adversarial learning strategy.

Due to the scarcity of medical images, it is usually necessary to use some data enhancement methods to expand the dataset. The generative models such as Variational AutoEncoder (VAE) [20], generative adversarial network (GAN) [21] are used. GAN is one of the most popular methods at present. It generates new data that conforms to the real data distribution through the train process of a discriminator and a generator. Many researchers constantly improve the generation effect of GAN, such as Gulrajani et al. [22] designed WGAN using wasserstein distance instead of cross entropy loss, and Mirza et al. [23] input conditional information and hidden variables into the generator and discriminator to control the type of generated data. After that, Brock et al. [24] designed BigGAN which slice the conditional information and hidden variables into different stages of the generator to achieve data generation under specific conditions, and improved the generation effect by increasing the batch size. Mariani et al. [25] proposed BAGAN, which used a VAE to obtain the distribution of data of different categories, and used this input generator to generate data of specified categories. However, due to the difficulty of training and the limitation of video storage, 3D GAN method is still relatively rare and is mostly used for the generation of brain cross modal images. There is less GAN model to generate CT images to balance data categories for the diagnosis of COVID-19.

In this study, we embed graph convolution aggregation and multi-scale self-attention fusion into a CNN backbone for global information optimization (GIONet) by enhancing the ability of global feature extraction. Experimental results on mixed-center task shows that our enhanced 3D CNN achieves quite promising results for COVID-19 diagnosis, which achieves an accuracy of 99.17%. In addition, 86.73% and 89.61% accuracy was obtained in the two tasks of cross center research. In summary, this work has three main contributions:

•We design a novel coronavirus disease generative adversarial network (COVID-GAN) to generate high-quality 3D CT data for imbalanced learning and COVID-19 diagnosis.

•We adopt a graph enhanced aggregation unit (GEAU) and multi-scale self-attention fusion unit (MSFU) in 3D CNN to enhance the ability of global features extraction, which can effectively extract features of lesions that are widely distributed in the lungs, such as GGO.

•We propose a domain adversarial strategy (DAS) for cross-center task to mitigate the difference of CT image characteristics between different centers.

2. Methods

2.1. Backbone

In this study, we have designed a COVID-GAN to balance the data before training. After obtaining the balanced data, we constructed a global information optimized network based on 3D SqueezeNet [26] due to its lightweight and high accuracy, making it easy to deploy and suitable for COVID-19 diagnosis. To enhance the global feature extraction capability of CNNs, we added two parallel branches in the middle of the backbone network. These branches are GEAU and MSFU, consisting of graph-enhanced aggregation operations and multi-scale self-attention fusion, respectively. GEAU utilizes the relationship between graph structure nodes and edges to enhance feature extraction, while MSFU utilizes self-attention to capture long-distance dependencies. Once feature extraction is complete, the input image is classified into its respective domain and disease categories using a domain classifier and a disease classifier. The overall block diagram of our approach is depicted in Fig. 1.

Fig. 1.

Global information optimized network and domain adversarial strategy. GEAU represents the graph enhanced aggregation unit and MSFU represents the multi-scale self-attention fusion unit.

2.2. COVID-GAN for imbalanced-learning

Based on our statistics and analysis of the dataset, we can clearly see that there is an imbalance in the categories of the dataset. As shown in Table 1, in the training set of the mixed-center task, the number of cases in the healthy control group (normal) is 516, the number of cases in CAP is 810, and the number of cases in COVID-19 is 4466. In addition, in the training sets of the two cross-center tasks, the distribution of the number of samples of different categories in the dataset is 749:1007:4097 and 378:942:1029. Fig. 2 shows the imbalance in data distribution clearly and intuitively.

Table 1.

The detailed distribution of datasets in different tasks.

Task	Train			Val			Test
	Normal	CAP	COVID-19	Normal	CAP	COVID-19	Normal	CAP	COVID-19
Mixed-center	516	810	4466	220	258	281	218	208	293
Task1	749	1007	4097	121	171	708	84	98	235
Task2	378	942	1029	93	236	258	84	98	235

Fig. 2.

The distribution of data of different categories in mix center task.

Such severely unbalanced data sets will seriously affect the COVID-19 diagnosis performance. Therefore, we have designed a new conditional generation countermeasure network adapted to 3D data, which expand the sample data of a small number of categories and achieve the purpose of data balance. Accordingly, the final diagnosis performance of COVID-19 can be enhanced.

The specific implementation idea of COVID-GAN is shown in Fig. 3, which is composed of two parts: generator and discriminator. Among them, the input of the generator is a noise with a length of 480 that conforms to the Gaussian distribution, and a unique heat coding of an image category. First, the Gaussian noise is divided into six equal length sub-vectors, and the first sub-vector is expanded into a 1 × 4×4 3D matrix, which is the initial input hidden variable of the generator. After five stages of three bits convolution and up-sampling, a 32 × 128×128 generated image is finally obtained. More specifically, the remaining five Gaussian noise sub-vectors will be gradually substituted in different stages of the generator after being connected with the category coding. As the condition control information of the generator, they will control the category of the image generated by the generator and maintain the diversity of the generated image. In this process, the feature will be weighted using the condition control information coding and can be described in Eq (1):

$$f_{out}=\mathrm{BN}\left(f_{in}\right)\times [1+{\mathrm{Linear}}_1\left(z_{condition}\right)]+{\mathrm{Linear}}_2\left(z_{condition}\right)$$ (1)

where $\mathrm{BN}$ means batch normalization, $z_{condition}$ means the hidden variable with conditional information, ${\mathrm{Linear}}_1$, ${\mathrm{Linear}}_2$ means two different linear transformation.

Fig. 3.

COVID-GAN for COVID-19 diagnosis with imbalanced data. $z_{condition}$ means the hidden variable with conditional information.

In the discrimination process, the input image is a real or false image and the category coding of the image. The input image is gradually extracted as a feature vector through multi-layer convolution and down-sampling layers, and the final feature value is obtained by weighting the feature vector with the category coding for discrimination. It is worth mentioning that the self-attention calculation process is added to the generator and discriminator to enhance the global information of features. The loss function of the generated confrontation process is described in Eq (2):

$$L_{GAN}=\underset{G}{min}\underset{D}{max}{E}_{x\sim P_r}\left[logD\left(x\right)\right]+E_{\tilde{x}\sim P_g}\left[log(1-D\left(\tilde{x}\right))\right]$$ (2)

where $E_{x\sim P_r}$ means that $x$ is a sample from the real distribution and $E_{\tilde{x}\sim P_g}$ means the $\tilde{x}$ is a sample from the generated data.

Finally, the generated CT images will be resampled to 64 × 256 ×256 for training, which is consistent with the original image.

2.3. Graph enhanced aggregation unit

Recently, graph-based methods have gained popularity due to their high efficiency in relational reasoning. Following the approach of [27], which utilized graph convolution for global reasoning. We proposed graph enhanced aggregation unit which illustrated in Fig. 4. Firstly, we transformed the coordinate space into a hidden space by a 1 × 1 × 1 convolution layer. Next, we constructed a fully-connected adjacency matrix $V$ by a learnable projection function shown in Eq. (3). And then, the features was calculated by two graph convolutional layers and reversed to coordinate space. The total GEAU process can be represented by Eq. (4), where $f$ denotes the output features of the convolutional layer, $V$ denotes the fully-connected adjacency matrix, each node stores features, and $W_{g\text{\_}1}$, $W_{g\text{\_}2}$ denotes the parameter matrix used in graph convolution.

$$\mathit{V}=\mathrm{projection}\left(\mathrm{Conv}\left({\mathit{F}}_{\mathit{in}}\right)\right)\times \mathrm{Conv}\left({\mathit{F}}_{\mathit{in}}\right)$$ (3)

$$Outpu{t}_{GEAU}={\mathit{F}}_{\mathit{in}}+\mathrm{Conv}(\mathrm{reverse}\left(\mathrm{Conv}\left({\mathit{F}}_{\mathit{in}}\right)\right)\times {\mathit{W}}_{\mathit{g}\text{\_}\mathbf{1}}{\mathit{V\; W}}_{\mathit{g}\text{\_}\mathbf{2}})$$ (4)

Fig. 4.

Graph enhanced aggregation unit. ${\mathit{F}}_{\mathit{in}}$ represents the feature from each stage of the backbone and is also the input feature of this unit, GCN means the normal graph convolutional network and $Outpu{t}_{GEAU}$ means the output feature of GEAU. $\mathit{V}$ denotes the fully-connected adjacency matrix, each node stores features, and ${\mathit{W}}_{\mathit{g}\text{\_}\mathbf{1}}$, ${\mathit{W}}_{\mathit{g}\text{\_}\mathbf{2}}$ denotes the parameter matrix used in graph convolution.

2.4. Multi-scale self-attention fusion unit

In addition to the Graph enhanced aggregation unit used in the previous section to enhance the global information extraction capability of the model, we also use a multi-scale self-attention calculation method, which can strengthen the long-distance relationship between features and further optimize global information. Unlike traditional Transformers [28], inspired by the dilated convolution operation and Swin-Transformer [29], this paper uses three parallel multi-head self-attention calculation branches, which perform window-based self-attention calculation on features operated with different dilation factors (e.g., 0, 1, 2). This further expands the global perception ability by sampling a larger range of features without significantly increasing computational costs. The specific operation is shown in Fig. 5, where $Q$, $K$ and $V$ represent query, key and value vectors, respectively. After the self-attention calculation in each branch, the dimensions of the features are different due to the dilation sampling. Therefore, before feature fusion, the feature dimensions need to be restored by padding with zero elements. The padding interval is consistent with the previous dilation sampling coefficient. After unifying the dimensions, the fusion is achieved by element-wise addition. Additionally, layer normalization (LN) layer and multilayer perceptron (MLP) layer. The calculation process of MSFU can be described in Eq (5), where $df$ represents dilation factors and $d$ stands for dimension.

$$Outpu{t}_{MSFU}=\sum _{df=0}^2\mathrm{MLP}\left(\mathrm{LN}(\mathrm{Softmax}\left(\frac{Q{K}^{\mathrm{T}}}{\sqrt{d}}\right)\times V)\right)/3$$ (5)

Fig. 5.

Multi-scale self-attention fusion unit. ${\mathit{F}}_{\mathit{in}}$ represents the feature from each stage of the backbone and is also the input feature of this unit, Embed means the patch embed operation, MLP means the Multilayer Perceptron and $Outpu{t}_{MSFU}$ means the output feature of MSFU.

2.5. Domain adversarial strategy for cross-center task

Due to the varied equipment types and scanning methods utilized by different centers, the CT images exhibit distinct features that can pose a challenge for classifiers. In order to address this issue, we have employed a domain adversarial training approach in this study, with the aim of minimizing the feature discrepancies between centers. The schematic diagram of domain adversarial learning is shown in Fig. 6, where ${\ast }^s$ and ${\ast }^t$ represent the source domain and target domain respectively. $x^{\ast }$ and $f^{\ast }$ means input images and features. In addition, $y^{\ast }$ and $d^{\ast }$ represent the class labels and domain labels. $\tilde{y}$ and $\tilde{d}$ mean the predicted disease type and predicted domain. The strategy can be described as follow: images were input into the GIONet and gained the features. Then, a diseases classifier and a domain discriminator were adopted to produce the prediction. And last, two loss functions were used to optimize the distance between features in different domains while maximizing the distance between features in different diseases.

Fig. 6.

Schematic diagram of domain adversarial strategy. Forwardprop denotes the forward propagation and Backprop denotes the back propagation. Enhanced CNN means the backbone in Fig. 1.

The loss function used in our domain adversarial training strategy is shown in Eq (6), where $L_y$ represents the loss function of the disease classifier, $y^s$ represents the real label of the disease, ${\tilde{y}}^s$ represents the disease type predicted by the network, $L_d$ represents the loss function of the domain discriminator, $d^s$ and $d^t$ represents the domain label, ${\tilde{d}}_s$ and ${\tilde{d}}_t$ represents the domain predicted by the network, $N$ denotes the total number of samples, while $source$ refers to the source domain, which includes the training and validation sets. Conversely, $target$ represents the target domain, which is the test set. The target domain only provides its domain label, and the class label is not disclosed. The loss functions used are $L_y$ and $L_d$, both of which are cross-entropy.

$$L_{DAS}=\sum _{\genfrac{}{}{0ex}{}{i=0}{source}}^N{L}_y({\tilde{y}}^s,y^s)-\sum _{\genfrac{}{}{0ex}{}{i=0}{source}}^N{L}_d({\tilde{d}}^s,d^s)-\sum _{\genfrac{}{}{0ex}{}{i=0}{target}}^N{L}_d({\tilde{d}}^t,d^t)$$ (6)

3. Results

3.1. Data and preprocessing

In this study, we utilize the publicly available COVIDx CT-3 dataset [30], which contains imaging data of COVID-19 patients from multiple countries and centers. The dataset includes 7270 CT scans of 6068 patients, comprising 431025 2D slice images collected from 11 data centers across more than 17 countries. We transform the 2D slice data into a 3D voxel image, where we upsample or downsample to 64 in the $z$-axis and 256 in the other two directions due to the inconsistent number of slices in the axial direction. This results in a final 3D volume size of [64, 256, 256]. We perform experiments on various tasks using this large dataset, and the detailed data partitioning is presented in Table 1. The mixed-center task employs data from all centers and splits the dataset into training, validation, and testing sets according to COVIDx CT-3 [30]. In Task 1, we use two centers for testing and the remaining centers for training. Task 2 uses the same test set as Task 1 but reduces the number of centers used for training.

3.2. Implementation details

In this research, we trained our models in parallel on four TITAN RTX GPUs with 24 GB memory, utilizing the PyTorch framework. The training process involved an initial learning rate of $5\times 10^{-4}$, and a learning rate decay mechanism with 0.95 decay coefficient was employed. Our experiment will be conducted for 10,000 epochs and the best model will be obtained through an early stopping mechanism. In addition, we have set the batch size to 8. To assess the performance of our GIONet with COVID-GAN and DAS, we compared it with several state-of-the-art models, including ResNets [31], MobileNetV2 [32], ShuffleNet [33], EfficientNet [34], and SqueezeNet [26]. In addition, Transformer methods that have demonstrated outstanding performance in visual tasks, such as ViT [35] and Swin-Transformer [36], will also be compared in the experiment, using their 3D versions. While also considering prior research on the same dataset, such as NASNet-A-Mobile [37], COVID-Net CT L [38], and COVID-Net CT S [39]. The results of the above three methods in Table 1 are directly quoted from their reported results in their respective papers. We measured the performance of our models using various metrics, including accuracy (ACC), sensitivity (SEN), precision (PRE), specificity (SPE), F1-score, and area under the curve (AUC), which were calculated using the macro-average method. We evaluated our COVID-GAN using the Fréchet Inception Distance (FID) [40] and Kernel Maximum Mean Discrepancy (MMD) [41]. The calculation formulas for the above metrics are shown below:

$$ACC=\frac{1}{3}\sum _{i=0}^3\frac{T{P}_i+T{N}_i}{T{P}_i+T{N}_i+F{P}_i+F{N}_i}$$ (7)

$$SEN=\frac{1}{3}\sum _{i=0}^3\frac{T{P}_i}{T{P}_i+F{N}_i}$$ (8)

$$PRE=\frac{1}{3}\sum _{i=0}^3\frac{T{P}_i}{T{P}_i+F{P}_i}$$ (9)

$$SPE=\frac{1}{3}\sum _{i=0}^3\frac{T{N}_i}{T{N}_i+F{P}_i}$$ (10)

$$F1-score=2\frac{PRE\times SEN}{PRE+SEN}$$ (11)

$$AUC=\frac{1}{3}\sum _{i=0}^3\frac{\sum _{j\in positiv{e}_i}ran{k}_j-\frac{M_i(M_i+1)}{2}}{M_i\ast N_i}$$ (12)

$$FID={\Vert {\mu }_r-{\mu }_g\Vert }^2+Tr\left({\Sigma }_r+{\Sigma }_g-2{\left({\Sigma }_r{\Sigma }_g\right)}^{1/2}\right)$$ (13)

where $\mu$ is the empirical mean, $\Sigma$ is the empirical covariance, $Tr$ is the trace of a matrix, $r$ represents the real dataset, and $g$ represents the generated dataset. Specifically, the calculation process of FID is as follows: Generate n*2048 vectors using InceptionV3 for n images in the target images, take the average to obtain ${\mu }_r$. Generate m*2048 vectors using InceptionV3 for m images in the generated images, take the average to obtain ${\mu }_g$. Calculate ${\Sigma }_r$ and ${\Sigma }_g$ from ${\mu }_r$ and ${\mu }_g$, and finally obtain FID.

$$MMD={\mathbb{E}}_{x,x^{\prime }\sim P_r}\left[k(x,x^{\prime })\right]-2{\mathbb{E}}_{x\sim p_{data},y\sim P_g}\left[k(x,y)\right]+{\mathbb{E}}_{y,y^{\prime }\sim P_g}\left[k(y,y^{\prime })\right]$$ (14)

MMD used to measure the distance between two probability distributions. In the MMD formula, $P_r$ is the real data distribution, $P_g$ is the generated data distribution, $x$ is a sample drawn from the real data distribution, $y$ is a sample drawn from the generated data distribution and $k(x,x^{\prime })$ is the Gaussian kernel function used to measure the similarity between samples $x$ and $x^{\prime }$.

3.3. Ablation study

To clearly demonstrate the effectiveness of each module, we have conducted many ablation experiments. Table 2 shows the performance improvements of generated images, GEAU and MSFU for training. We can observe that designed module improves the diagnosis performance. COVID-GAN, GEAU and MSFU respectively increased the accuracy of the backbone network by 0.93, 0.18, and 0.37 percentage points. In addition, the combination of several modules has also contributed to different degrees of improvement in the model, and the complete model achieved the best performance.

Table 2.

The ablation study of different modules, Boldface denotes the best performance(%).

GAN	GEAU	MSFU	ACC	PRE	SEN	SPE	F1	AUC
			97.59	96.60	96.27	98.09	96.43	98.62
✓			98.52	97.84	97.81	98.85	97.82	99.15
	✓		97.77	96.73	96.68	98.27	96.71	99.19
		✓	97.96	97.06	96.91	98.40	96.98	98.81
✓	✓		98.78	98.15	98.25	99.09	98.20	99.25
✓		✓	98.69	97.99	98.09	99.02	98.04	99.32
	✓	✓	98.05	97.22	97.02	98.47	97.12	99.38
✓	✓	✓	99.17	98.77	98.79	99.36	98.78	99.67

3.4. Comparative results

Table 3 displays the performance of our COVID-GAN following academic conventions. It is evident that our COVID-GAN obtains the lowest values for both MMD and FID, indicating that our GAN is capable of generating diverse CT images that closely approximate the true distribution. Table 4 showcases a comparative evaluation of our method and several neural network approaches on mixed-center task data. We observe that the use of SqueezeNet as the backbone leads to superior outcomes in various comparison metrics. Our GIONet achieves the best overall accuracy of 99.17%, with sensitivity and specificity values of 98.79% and 99.36%, respectively. Our results outperform both conventional neural networks and those reported in prior studies. Additionally, we can visualize the classification results for the mixed-center task through Fig. 7, where our method has the highest number of samples in the diagonal direction of confusion matrix and achieves the best results.

Table 3.

The performance of our COVID-GAN. Boldface denotes the best performance(%).

Methods	MMD	FID
BAGAN	3.55	75.24
BigGAN	2.01	35.41
WGAN	3.23	68.52
CGAN	2.52	42.33
Ours	1.47	30.54

Table 4.

The performance comparison of different methods in mixed dataset. Boldface denotes the best performance(%).

Methods	ACC	PRE	SEN	SPE	F1	AUC
ResNet-18	96.48	95.03	94.62	97.22	94.83	98.71
ResNet-34	93.60	90.82	90.27	94.96	90.55	95.59
MobileNet	94.44	91.57	92.08	95.77	91.82	95.47
ShuffleNet	95.36	93.05	93.33	96.46	93.19	97.72
EfficientNet	85.91	90.46	78.57	89.02	79.50	96.05
SqueezeNet	97.59	96.60	96.27	98.09	96.43	98.62
ViT	96.75	95.17	95.15	97.50	95.16	96.88
Swin-Transformer	97.03	95.64	95.47	97.70	95.55	97.98
NASNet-A-Mobile [37]	98.8	–	98.7	–	–	–
COVID-Net CT L [38]	98.4	–	98.1	–	–	–
COVID-Net CT S [39]	98.3	–	97.3	–	–	–
Ours	99.17	98.77	98.79	99.36	98.78	99.67

Fig. 7.

The confusion matrix of mixed-centers task.

Furthermore, Table 5 presents comparative experimental results for two cross-center tasks. Our method, utilizing domain adversarial training strategies, attains the best performance with 86.73% accuracy in the first cross-center task and 89.61% accuracy in the second task, surpassing other approaches in all measures. The abbreviation “w/o DAS” indicates that our network does not incorporate domain adversarial strategy. It is worth noting that the experimental data of both Task 1 and Task 2 in our method are balanced by COVID-GAN. Finally, although our method still exhibits some weaknesses in distinguishing between CAP and COVID-19, it outperforms other methods, as evidenced by Fig. 8. The results of Task 1 were not as good as Task 2, despite Task 1 having more training data, because the data distribution in Task 2 was more balanced, and the lesions in Task 2 were more representative of the overall population. Additionally, Task 1 had a larger class imbalance, which made it more challenging for the model to learn and generalize well to the test set.

Table 5.

The performance comparison of different methods in cross center dataset. Boldface denotes the best performance(%).

Task	Methods	ACC	PRE	SEN	SPE	F1	AUC
Task1	ResNet-18	81.77	67.08	64.77	81.00	65.90	72.12
	ResNet-34	81.45	70.60	64.23	80.98	67.26	75.77
	MobileNet	81.45	67.30	63.69	80.40	65.44	72.99
	ShuffleNet	81.93	75.40	64.11	80.96	69.30	74.39
	EfficientNet	81.14	65.54	62.24	80.46	63.85	74.07
	SqueezeNet	83.85	74.54	68.20	82.95	71.22	79.03
	ViT	80.02	58.65	62.81	79.69	60.66	70.25
	Swin-Transformer	82.41	69.23	65.33	81.40	67.23	80.52
	Ours(w/o DAS)	85.93	79.61	73.07	85.41	76.20	80.85
	Ours	86.73	85.98	73.61	85.87	79.32	82.32
Task1	ResNet-18	76.98	63.94	68.78	81.06	66.27	80.40
	ResNet-34	74.42	61.98	66.26	79.53	64.05	79.92
	MobileNet	84.33	76.05	76.50	86.71	76.27	86.22
	ShuffleNet	77.62	62.86	66.97	80.78	64.85	79.67
	EfficientNet	57.31	36.15	54.17	71.72	43.36	70.14
	SqueezeNet	85.45	77.25	77.49	86.98	77.37	87.52
	ViT	84.81	76.38	70.72	84.35	73.44	85.08
	Swin-Transformer	85.29	77.51	71.74	84.90	74.51	88.20
	Ours(w/o DAS)	88.01	91.94	74.38	86.26	82.23	89.15
	Ours	89.61	92.78	77.83	88.10	84.65	90.21

Fig. 8.

The confusion matrix of Task1 and Task2.

From Table 6, we can see that our method has 3.45M parameters and a computational complexity of 23.2GMac, while SqueezeNet, which has the second-best performance, has 1.84M parameters and a computational complexity of 15.61. With a small increase in the number of parameters, we can achieve an improvement in accuracy, and in our experimental environment, such an increase in parameters is completely acceptable.

Table 6.

The computational complexity comparison of different methods.

Methods	Flops(GMac)	Params(M)
ResNet-18	126.04	33.16
ResNet-34	216.72	63.47
MobileNet	9.75	2.36
ShuffleNet	2.63	1.30
EfficientNet	5.55	4.69
SqueezeNet	15.61	1.84
ViT	10.79	84.07
Swin-Transformer	103.1	27.54
Ours	23.2	3.45

3.5. Visualization

We have visualized the generated 3D CT image from the designed generator from the transverse section as shown in Fig. 10. It shows that the quality of the image we generated is good. From the last line of Fig. 10, the red boxes indicate the difference from Normal, CAP and COVID-19. There are shadows in the lungs of CAP cases and GGO in the lungs of COVID-19 cases. However, there are some problems with the images generated by other methods. For example, BAGAN’s images generated by mode collapse are very similar, or CGAN produces some strange textures. Our method can obtain images with similar structure but different categories by changing the conditional information under the fixed hidden variable input.

Fig. 10.

The 2D visualization of generated volume CT images from different GAN methods and the last line are the results of our COVID-GAN.

To visually demonstrate the effectiveness of the feature extraction and classification methods proposed in our work, we applied t-distributed stochastic neighbor embedding (t-SNE) visualization to the output features of each method [42]. Fig. 11 shows that our method can effectively distinguish different disease features. In contrast, the features generated by other methods often have a large overlap, while our method separates the features of different categories further apart. Additionally, domain adversarial training can minimize the distance between features from different domains and maximize the distance between features from different diseases. This can be observed in Fig. 12, where domain adversarial training effectively reduces the feature distance between different domains and increases the feature distance between different disease categories.

Fig. 11.

The visualization of t-SNE characteristics of different methods in mixed-center task, in which green represents normal, blue represents CAP, and red represents COVID-19.

Fig. 12.

The visualization of t-SNE characteristics in cross-center tasks, in which blue dots represent the characteristics of the source domain, red dots represent the characteristics of the target domain, and different shapes represent different disease types.

From Fig. 9, it can be seen intuitively that the classification basis of our method mainly lies in some special features of the lungs, such as shadow, ground-glass opacity and other areas, while the classification basis of other methods is obviously incorrect. For example, in the case of ResNet34’s activation map, the classification basis is too broad to explain the differences in the model’s classification judgments for different categories of images.

Fig. 9.

The class activation map of different methods on three categories of data.

4. Conclusions

In summary, this paper presents a novel approach to enhance the global feature extraction ability of neural networks for the accurate classification of COVID-19, by leveraging graph enhanced aggregation unit and multi-scale self-attention fusion unit. Furthermore, a COVID-GAN is introduced to generate high-quality 3D CT images for data balancing learning, while domain adversarial training is utilized to optimize the performance of cross-center tasks by reducing the feature difference between domains. Experimental results demonstrate that the proposed method achieves exceptional accuracy, with 99.17% on mixed-center task and 86.73%, 89.61% accuracy on two cross-center tasks, surpassing existing methods in terms of feature extraction ability and classification performance. These findings suggest that the proposed approach holds great potential for accurate COVID-19 diagnosis.

5. Discussion

This work has several potential limitations. First of all, because our three-dimensional data comes from the reconstruction of two-dimensional slices, and the number of slices in some cases is very small, we need to repeat the sampling of some slices for many times during the reconstruction. Therefore, this part of the cases will produce distortion in the three-dimensional data, which will affect the experimental results. In the future, we will use the GAN method to reconstruct 3D images on a small number of slices of this part of data with a small number of slices to replace repeated sampling. Secondly, because there are only two types or even one type of data in this dataset on several centers, there are fewer centers that can be selected for training when conducting cross center dataset experiments. In the future, we will collect more multi center data to more widely verify our multi center domain adversarial COVID-19 diagnosis method. Finally, our COVID-GAN method can only expand the category balance data, and we hope that it can be improved in the future to achieve category missing completion, as well as domain confrontation adaptation for multi-centers. Finally, this work can also be extended to the diagnosis of other lung diseases in the future, such as nodules, inflammations, and so on.

Declaration of Competing Interest

The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.