Localization of lung abnormalities on chest X-rays using self-supervised equivariant attention

Abstract

Chest X-Ray (CXR) images provide most anatomical details and the abnormalities on a 2D plane. Therefore, a 2D view of the 3D anatomy is sometimes sufficient for the initial diagnosis. However, close to fourteen commonly occurring diseases are sometimes difficult to identify by visually inspecting the images. Therefore, there is a drift toward developing computer-aided assistive systems to help radiologists. This paper proposes a deep learning model for the classification and localization of chest diseases by using image-level annotations. The model consists of a modified Resnet50 backbone for extracting feature corpus from the images, a classifier, and a pixel correlation module (PCM). During PCM training, the network is a weight-shared siamese architecture where the first branch applies the affine transform to the image before feeding to the network, while the second applies the same transform to the network output. The method was evaluated on CXR from the clinical center in the ratio of 70:20 for training and testing. The model was developed and tested using the cloud computing platform Google Colaboratory (NVidia Tesla P100 GPU, 16 GB of RAM). A radiologist subjectively validated the results. Our model trained with the configurations mentioned in this paper outperformed benchmark results.

Supplementary information: The online version contains supplementary material available at 10.1007/s13534-022-00249-5.

Keywords: CAM; Pixel correlation module; ResNet50; Self-attention; Self-supervised equivariant attention; Siamese network; Weak supervision.

Fig. 1

Lung diseases and their appearance on anterior to posterior CXR pose a great challenge to validate in a CAD system (images from [1]). Few examples, a Atelatasis, b cardiomegaly, c effusion, d Infiltration

Fig. 2

a Multi-resolution ResNet backbone architecture that sequentially reintegrates spatial information from previous layer outputs using the Merge Block. b shows the internal structure of the Merge Block which uses upsampling and concatenation to merge feature maps.

Fig. 3

Structure of PCM. H, W, $C_{\mathit{in}}$ /$C$ denote the height, width and channel length of feature maps and original cams

Fig. 4

Siamese Network for SEAM during PCM training. PCM refines the pixel-wise CAM predictions by using context information from feature maps produced by the backbone feature extractor. Image augmentations by affine transformations create inconsistencies in the generated CAMs compared to the original CAMs. Equivariant regularization enables the use of these inconsistencies for PCM training. MSML loss provides an additional source of supervision

Fig. 5

a Shows the training loss curves during training for all models. b Shows the validation loss curves during training for all models.

Fig. 6

ROC Curve for the performance of MultiResolution-ResNet50

Fig. 7

Sample localization heatmaps and bounding boxes generated by standard ResNet50 (left), MultiRes-ResNet50 (Middle), and MultiRes-ResNet50 CAMs refined using PCM (right). In each figure, green bounding boxes indicate ground truths, whereas red bounding boxes indicate predicted predictions by the corresponding model