Synergistic Learning of Lung Lobe Segmentation and Hierarchical Multi-Instance Classification for Automated Severity Assessment of COVID-19 in CT Images

Segmentation of lung or lung lobe has been used as a common pre-requisite procedure for automatic diagnosis of COVID-19 based on chest CT images. Several deep learning methods have been proposed for the segmentation of lung in CT images with COVID-19. For instance, U-Net et al .[13] has been widely used for segmentation of both lung regions and lung lesions in COVID-19 applications [14, 15, 16, 7]. Qi et al . [7] use U-Net to delineate the lesions in the lung and extract radiometric features of COVID-19 patients with the initial seeds given by a radiologist for predicting hospital stay. Also, several variants of U-Net have been applied to the diagnosis or severity assessment of COVID-19. Jin et al . [6] design a two-stage pipeline to screen COVID-19 in CT images, and they utilize U-Net++ [17] to detect the whole lung region and to separate lesions from lung regions. Besides, V-Net [18] is also used in various segmentation applications. Shan et al . [5] integrates human-in-the-loop strategy into the training process of VB-Net (a variant of V-Net). The human-aided strategy is an intuitive way to address the issue of lacking manual labels during segmentation in CT images.

Both X-rays [19] and CT images [10] can provide effective information for the computer-assisted diagnosis of COVID-19. Compared with X-rays, chest CT imaging contains hundreds of slices, which is clearer and more precise but has to take more time for specialists to diagnose. Therefore, there is a great demand to use CT images for automated diagnosis of COVID-19. In general, the existing methods for COVID-19 diagnosis based on CT images can be roughly divided into two categories: 1) classification; 2) severity assessment. In the former category, many studies have been conducted to determine whether patients are infected with COVID-19 disease. For example, Chen et al . [2] exploits a UNet++ based segmentation model to segment COVID-19 related lesions in chest CT images of 51 COVID-19 patients and 55 patients with other diseases, and finally determine the label (COVID-19 or non-COVID-19) of each image based on the segmented lesions. Ying et al . [3] propose a CT diagnosis system, namely DeepPneumonia, which is based on the ResNet50 model to identify patients with COVID-19 from bacteria pneumonia patients and healthy people. In the second category, Tang et al .[4] proposed to first adopt VB-Net to separate the lung into anatomical sub-regions, and then use these sub-regions to compute quantitative features for training a random forest (RF) model for COVID-19 severity assessment (with labels of being non-severe or severe).

The scenario of multi-instance learning (MIL) [20, 21, 22] or learning from weakly annotated data [23] arises when only a general statement of the category is given, but multiple instances can be observed. MIL aims to learn a model that can predict the label of a bag accurately, and many recent studies have focused on implementing MIL via deep neural networks. For instance, Oquab et al . [23] train a deep model with multiple image patches of multiple scales as input, and aggregate the prediction results of multiple inputs by using a max-pooling operation. Besides, many studies [24, 25, 26] propose to formulate image classification as a MIL problem so as to address the weakly supervised problem. Moreover, MIL is particularly suitable for problems with only a limited number (e.g., tens or hundreds) of training samples in various medical image-based applications, such as computer-assisted disease diagnosis [27, 28, 29, 30]. For instance, Yan et al .[29] propose a two-stage deep MIL method to find discriminative local anatomies, where the first-stage convolutional neural network (CNN) is learned in a MIL fashion to locate discriminative image patches and the second-stage CNN is boosted using those selected patches. More recently, a landmark-based deep MIL framework [30] is developed to learn both local and global representations of MRI for automated brain disease diagnosis, leading to a new direction for handling limited training samples in the domain of medical image analysis. Since there are only a limited number of cases at hand, it is desirable to employ the multi-instance learning strategy for severity assessment of COVID-19 patients in chest CT images.

The framework of the proposed method is illustrated in Fig. 3, where the input is the raw 3D CT image and the output is the lung segmentation and severity assessment of COVID-19 patients (i.e., severe or non-severe). Specifically, each 3D CT image is processed via several image pre-processing steps. Then, a set of 2D image patches is randomly cropped from the processed image to construct an instance bag, and each bag represents a specific input CT image. This bag is regarded as the input of the proposed multi-task multi-instance U-Net (M²UNet). The M²UNet is designed to learn two tasks jointly, i.e., severity assessment of a COVID-19 patient and segmentation of the lung lobe.

As shown in Fig. 3, in M²UNet, an encoding module is first used for patch-level feature extraction of 2D patches in each input bag, followed by two sub-networks for joint severity assessment and lung lobe segmentation. Specifically, in the classification sub-network, these extracted patch-level features are fed into a feature embedding module and an image-level feature learning module to capture the local-to-global volume representation of the input CT image. With the learned volume features, a classification layer is finally used to assess the severity of each COVID-19 patient (i.e., severe or non-severe). In the segmentation sub-network, those patch-level features are fed into a decoding module to perform lung lobe segmentation for each patch in the input bag. Since these two sub-networks are trained jointly with a shared patch-level encoder, the context information provided by the segmentation results can be implicitly employed to improve the performance of severity assessment.

To eliminate the effect of the background noise in each raw 3D CT image, we crop each scan to only keep the region containing the human body, by using a threshold-based processing method. Specifically, we first binarize the image using the threshold of zero, through which the human tissues and the gas regions will be separated. Then, the human body region is cropped according to the binary mask. Each body region image has a size of at least $256\times 256$ for the axial plane in this work.

While image resampling is commonly used in many deep learning methods for segmentation and classification [31, 32], we do not resample the raw CT images in order to preserve their original data distributions. Since our method is clinical-oriented with inconsistent imaging qualities, CT images used in this study are not as clean as those in benchmark datasets [33]. For example, the physical spacing of our data has large variation, e.g., from $0.6\,mm$ to $10\,mm$ between slices, because of the use of different CT scanners and scanning parameters. Using a common interpolation method (e.g., tri-linear interpolation) to resample a CT image into $1\,mm$, one will introduce heavy artifacts to the image. Besides, only a few infection regions in each CT image are related to severity assessment. To this end, we employ the weakly supervised multi-instance learning (MIL) strategy for handing these inconsistent CT images. Specifically, for each pre-processed CT image, we randomly crop a set of 2D patches sampled from 2D slices (with each patch from a specific slice) in each image to construct an instance bag, and each bag is used to represent a specific CT image and treated as the input of the subsequent M²UNet. In this way, the inter-slice/patch relationships can be implicitly captured by our M²UNet. In addition, this MIL strategy represents each 3D image through a set of 2D image patches rather than sequential slices. This can partially alleviate the problem of data inconsistency, so our method has high practical value in real-world applications.

As shown in Fig. 3, using each bag (consisting of a set of 2D image patches) as the input data, the proposed M²UNet first employs an encoding module for patch-level feature extraction. Based on these features, the classification and segmentation sub-networks are then used to jointly perform two tasks, respectively, i.e., 1) severity assessment of the patients, and 2) segmentation of lung lobes in each patch. Specifically, the classification sub-network uses a unique hierarchical MIL strategy to extract the local-to-global representation of each input image, with an embedding-level MIL module, an image-level MIL module, and a classification module. The segmentation sub-network contains a decoding module to segment lung lobes of 2D image patches in each bag.

The detailed network architecture is listed in Table I. The combination of the encoder and decoder is U-Net like, with four down-sampling blocks in the encoder and four up-sampling blocks in the decoder. The outputs of the same level blocks in the encoder and decoder are concatenated and fed into the next block of the decoder. Limited by computational resources, all the convolutional layers in the encoder and decoder have the same number (i.e., $64$) of kernels, except the last block in the encoder. The last block of encoder outputs $512$ dimensional features to help build a more robust classification for severity assessment. The decoder outputs the corresponding segmentation mask of five types of lung lobes for each image patch.

While infection regions of the lung, related to COVID-19 (e.g., nodule and GGO) are usually located in regions of the CT image, the category of each CT image is labeled at the entire image level, rather than the region-level. That is, many regions are actually unrelated to the classification task for severity assessment.

Multi-instance learning (MIL) provides a useful tool to solve such a weakly supervised problem. Conventional MIL represents a 2D image as a bag, and each bag consists of multiple regions of the input image (i.e., instances). Their overall prediction is made at the bag-level by roughly two kinds of methods, i.e., the embedding-level methods and the instance-level methods. The former learns the relationship among instances by projecting them into a new embedding space. The latter directly generates the bag-level predictions by performing voting on the instance predictions. However, both methods are inapplicable for the classification of 3D images, as 3D images contain multiple 2D image slices, and the class labels are also related to some local regions of the slices.

In this work, based on the previous study on pathological images [12], we propose a hierarchical MIL strategy in the classification sub-network of our M²UNet to perform severity assessment of COVID-19 in 3D CT images, as shown in Fig. 3. As mentioned in Section III-B, we represent each input 3D volumetric image as a bag consisting of a set of 2D image patches, and these patches are regarded as the instances in the MIL problem settings. Formally, we first construct a bag with $n$ 2D patches cropped from the regional slices to represent each input CT image. Denote the $i$-th and the $j$-th 3D CT image as ${\mathcal{X}_{i}}$ and ${\mathcal{X}_{j}}$, respectively, where $\mathcal{X}_{i}=\{\phi_{i1}^{ins},\phi_{i2}^{ins},\cdots,\phi_{in_{i}}^{ins}\}$ and $\mathcal{X}_{j}=\{\phi_{j1}^{ins},\phi_{j2}^{ins},\cdots,\phi_{jn_{j}}^{ins}\}$. Here, $\phi_{kl}^{ins}\in\mathrm{R}^{d}$ $(k=1,2,\cdots,n_{k})$ indicates the $l$-th instance of the $k$-th image. Then, these 2D patches (size: $height\times width$) are fed into the encoding module for patch/instance-level feature extraction. These instance-level features are further fed into an embedding-level MIL module, which will be introduced later. After obtaining the instance-level features, the bag/image-level feature $\Phi_{i}$ are then generated by our proposed global contrast pooling (GCP) layer in the image-level MIL module.

As illustrated in Fig. 4, the proposed GCP layer aims to make the instance features closer to the relevant concepts, and also push those irrelevant instance features and concepts away from each other. In this work, the term ‘concept’ denotes the to-be-learned feature of GCP layer that is discriminative for severity assessment. Theoretically, the concept is a normalized weight to map features in instance feature space to an ordered embedding space. Specifically, in the GCP layer, we assume the bag-level feature $\Phi_{i}$ is represented by the relationship between instance features and $p$ concepts. Here, these concepts are learned to reveal the data structure in a global perspective. The bag-level feature is then denoted as a $p$ dimensional feature vector, with each dimension denoting the maximum similarity between one concept and all instance features. We use the cosine function to measure such relationships. Thus, the bag feature and the similarity can be written as

where $s_{im}$ ($m=1,\cdots,p$) is the maximum similarity between the instance features of the $i$-th bag and the $m$-th image-level concept $w_{m}$. $R(\cdot)$ denotes the commonly used regularization term used in deep networks. With Eqs. (1)-(2), one can observe that the proposed GCP layer can automatically learn the concepts that are related to those discriminative instances, thus reducing the influence of those irrelevant instances. Note such a GCP layer can be also used in other weakly supervised problems, where only a small portion of regions in an image are related to the task at hand (such as MRI-based brain disorder diagnosis [30]).

We further use an embedding-level MIL module (with a GCP layer) to learn embedding-level representations, by regarding each image patch as a bag and the intermediate patch-level features produced by the encoder as instances. In this way, the relationships among small regions in each patch can be modeled by our method. Based on the embedding-level features, an image-level MIL module (with a GCP layer) is further used to generate the volume features. Based on the volume features, we use a fully-connected layer followed by a cross-entropy loss to predict the severity score (i.e., severe or non-severe) of each input CT image. The final loss function in the proposed hierarchical MIL network for severity assessment can be formulated as

where $fc(\cdot)$ denotes the mapping function of the fully-connected layer, and $y$ denotes the severity type confirmed by clinicians.

The segmentation task is supervised by the aggregation of cross-entropy loss and Dice loss as follows

where $\hat{p}_{n}^{c}$ and $l_{n}^{c}$ denote the predicted and ground-truth segmentation masks for the $n$-th patch in the $c$-th category. In this work, we segment $c=7$ categories, including five parts of lung lobes and the background. It is worth noting that most of the cases in our dataset do not have ground-truth segmentation masks. For these cases, we simply avoid calculating the segmentation loss for them.

Finally, the losses in Eqs. 3 and 4 are trained simultaneously in a multi-task learning manner, and the overall loss of the proposed method is written as

where $\lambda$ is the trade-off parameter used to balance the contributions of these two tasks. In this work, $\lambda$ is empirically set to $0.01$.

The proposed method is implemented based on the open-source deep learning library Pytorch. The training of the network is accelerated by four NVidia Tesla V100 GPUs (each with 32 GB memory). For feasible learning of the lung region images, we clamp the intensities of the image into $[-1200,0]$, which indicates that we use the width of $1200$ and the level of $-600$ for the pulmonary window. Then, the data is normalized to the value of $[0,255]$, as used by other deep learning methods. The dataset is highly imbalanced as the number of severe patients is much fewer than the non-severe patient. The ratio of the severe patient is less than $20\%$ in our dataset. Therefore, we augmented the data by directly duplicated the severe cases in the training set. This can be done because the proposed method uses a random cropping strategy to construct the inputs. This makes the duplicated cases not the same to each other for the training of the network. We also use the random cropping strategy in the testing stage, by assuming that the data distribution is already well learned in training. Other cropping strategies, e.g., center cropping, may not be suitable here, as the center of pulmonary is dominated by the trachea and other tissues.

In both training and testing stage, we randomly crop $200$ image patches from each input 3D CT image to construct the image-level bag (i.e., with the bag size of $n=200$). And we use the output of the encoder to construct the embedding-level bag that contains $8\times 8$ feature maps. We train M²UNet using the learning rate of $0.01$ with a decay strategy of ‘Poly’ (with the power of $0.75$). The network is optimized by a standard Stochastic Gradient Descent (SGD) algorithm with $100$ epochs. And the weights are decayed by a rate of $1\times 10^{-4}$ with the momentum of $0.9$.

The real COVID-19 dataset contains a total of $666$ 3D chest CT scans acquired from $242$ patients who are confirmed with COVID-19 (i.e., RT-PCR Test Positive). These CT images are collected from seven hospitals with a variety of CT scanners, including Philips (Ingenuity CT iDOSE4), GE (Bright speed S), Siemens (Somatom perspective), Hitachi (ECLOS), and Anke (ANATOM 16HD). The images are of large variation in terms of the image size of $512\times(512\sim 666)\times(23\sim 732)$, and the spatial resolution of $0.586\sim 0.984\,mm$, $0.586\sim 0.984\,mm$ and $0.399\sim 10\,mm$. Obviously, diagnosis based on these images is a very challenging task. The severity of the patient is confirmed by clinicians, following the guideline of 2019-nCoV (trail version 7) issued by the China National Health Commission. The severity of the patient is categorized into four types, i.e., mild, moderate, severe, and critical. In clinical practice, patients are often divided into two groups with different treatment regimens, i.e., severe and non-severe. The segmentation of 152 out of 666 images were delineated by an AI-based software and confirmed by experienced radiologists. In this work, we employ this partitioning strategy. That is, mild and moderate are treated as non-severe, while severe and critical are regarded as severe. Therefore, the task of severe assessment is formulated into a binary classification problem. Therefore, the dataset is partitioned into $51$ severe and $191$ non-severe patients.

We first compare the proposed M²UNet with four state-of-the-art methods in [4, 36, 34, 35] for severity assessment of COVID-19 patients. The first two methods [4, 36] are both based on hand-crafted features of CT images, while the last two [34, 35] are deep learning-based methods that can learn imaging features automatically from data. Specifically, Tang et al . [4] first segment the lung, lung lobe and lesions in CT images. Then, the quantitative features of COVID-19 patients, e.g., the infection volume and the ratio of the whole lung, are calculated based on the segmentation results. The prediction is done by a random forest method. Yang et al . [36] propose to aggregate infection scores calculated on 20 lung regions for severity assessment. Besides, the ResNet50+Max method [34] is also compared for patch-wise classification. ResNet50+Max is a non-MIL method, which has a ResNet-50 network architecture and performs image-level classification through max-voting. In addition, we apply the Gated Att. MIL method proposed in [35] on our dataset, which is a one-stage MIL method with an attention mechanism. For fair comparison, this method shares the same multi-instance pool as our M²UNet.

We further compare our method with two state-of-the-art methods for lung lobe segmentation, including 1) UNet [13], and 2) UNet++ [17]. The parameter settings for these five competing methods are the same as those in their respective papers.

To evaluate the influence of the proposed multi-task learning and hierarchical MIL strategies used in M²UNet, we further compare M²UNet with its two variants: 1) M²UNet with only the classification sub-network (denoted as Clas. Only), 2) M²UNet with only the segmentation sub-network (denoted as Seg. Only).

A five-fold cross-validation (CV) strategy is used in the experiments for performance evaluation. Specifically, the whole dataset is first randomly partitioned into five subsets (with approximately equal sample size of subjects). We treat one subset as the testing set ($20\%$), while the remaining four subsets are combined to construct the training set ($70\%$) and validation set ($10\%$). The validation set here is used for selecting the hyper-parameters. This process is iterated until each subsets serve as a testing set once. The final results are reported on the test set.

Two tasks are included in the proposed method, i.e., classification of severity assessment, and 2) segmentation of the lung lobe. For performance evaluation, two sets of metrics are used in these two tasks, with the details given below.

We further evaluate the influence of two major strategies used in our M²UNet, i.e., 1) the hierarchical MIL strategy for classification, and 2) the multi-task learning strategy for joint severity assessment and lung lobe segmentation.

We further investigate the performance of our method using different size of bags, and the results are shown in Table VI. Specifically, we vary the bag size within $\{50,80,100,150,200\}$. As shown in the table, the performance of M²UNet for classification changes along with the bag size. The proposed method achieves the best performance with the bag size of $200$. Another observation is that, the table suggests the performance of the proposed AI-based severity assessment model is not sensitive with the bag size when larger than $100$, indicating that at least $100$ patches are required for the proposed method to achieve an acceptable result.