A novel multiple instance learning framework for COVID-19 severity assessment via data augmentation and self-supervised learning

Abstract

How to fast and accurately assess the severity level of COVID-19 is an essential problem, when millions of people are suffering from the pandemic around the world. Currently, the chest CT is regarded as a popular and informative imaging tool for COVID-19 diagnosis. However, we observe that there are two issues - weak annotation and insufficient data that may obstruct automatic COVID-19 severity assessment with CT images. To address these challenges, we propose a novel three-component method, i.e., 1) a deep multiple instance learning component with instance-level attention to jointly classify the bag and also weigh the instances, 2) a bag-level data augmentation component to generate virtual bags by reorganizing high confidential instances, and 3) a self-supervised pretext component to aid the learning process. We have systematically evaluated our method on the CT images of 229 COVID-19 cases, including 50 severe and 179 non-severe cases. Our method could obtain an average accuracy of 95.8%, with 93.6% sensitivity and 96.4% specificity, which outperformed previous works.

Keywords: COVID-19; Chest CT; Data augmentation; Multiple instance learning; Self-supervised learning.

Fig. 1

Examples of chest CT images with severe infection (left) and non-severe infection (right) of COVID-19. The yellow arrows indicate representative infection regions. (For interpretation of the references to colour in this figure legend, the reader is referred to the web version of this article.)

Fig. 2

An overview of our method. In our method, the CT images are cropped into patches, which are then packed into MIL bags. In the $k{}^{th}$ epoch of training process, the data for supervised training consists of real bags (i.e., training CT images) and virtual bags generated in the ${(k-1)}^{th}$ epoch. Besides, real bags are also used for the auxiliary self-supervised learning task (while virtual bags are not). After the training stage, the trained MIL model will take the testing CT images (also modeled as MIL bags) as input to predict their labels (i.e., severe or non-severe).

Fig. 3

A brief illustration of the notions of CT slices, instances (patches) and bags. A CT image contains CT slices, and the slices are cropped to non-overlapping patches, which are considered as instances. The patches from the same CT image make up a MIL bag, with a bag-level label “severe” or “non-severe”.

Fig. 4

The framework of deep MIL model. Firstly, the instance features are extracted. Secondly, the attention weights of the instance features are determined by the network. Then, the MIL pooling layer combines the instance features to generate a bag feature. Finally, the bag feature is mapped by a fully connected (FC) layer to decide the label.

Fig. 5

An example of key instances in a positive bag. It indicates that the patches are likely to be related to the severe infection regions. Note that, we rescaled the attention weights of the patches in the same slice using $a_k^{{}^{\prime }}=a_k/{\sum }_i{a}_i$.

Fig. 6

A sketch map of generating virtual bags. For positive bags, instances with high weights are appended to the list of key instances while instances with low weights to the list of regular instances. Virtual bags are generated by randomly sampling key instances and regular instances.

Fig. 7

A CT slice is divided to 12 patches. For the relative patch location task, we are able to create 16 pairs of patches: $(a_0,a_1),\cdots ,(a_0,a_8)$ and $(b_0,b_1),\cdots ,(b_0,b_8)$. For the absolute patch location task, we directly predict each patch’s location among $l_1,\cdots l_{12}$.

Fig. 8

The age distribution of the patients in our dataset.

Fig. 9

Visualization of the bag-level features extracted in different configurations. The left corresponds to (A), while the right corresponds to (D). Red and green points stand for severe and non-severe cases, respectively. (For interpretation of the references to colour in this figure legend, the reader is referred to the web version of this article.)

Fig. 10

Visualization of the attention mechanism in our method. Presented above are some examples of the (slice-wise rescaled) attention weights of patches. Severe infection regions identified by experts are marked with yellow boxes. It can be seen that the patches with high weights are probably relevant to severe infection. (For interpretation of the references to colour in this figure legend, the reader is referred to the web version of this article.)

Fig. 11

Performance of different $\alpha$. The best $\alpha$ is 0.025.

Fig. 12

Performance of different $\mu$. The best $\mu$ is 0.3.