A Structure-Aware Relation Network for Thoracic Diseases Detection and Segmentation

Most existing works on chest X-rays focus on disease classification and weakly supervised localization[30, 21, 8, 20, 35, 34], due to the fact that most current public datasets contain only class labels or very limited amounts of box annotations. CheXpert[12] proposed by Irvin et al. contains 224,316 frontal and lateral view chest X-ray images. In training set, the labels (1/0/uncertain of the presence of 14 classes of observations) are extracted from radiology reports by a carefully designed labeler. In validation and test datasets, the labels are labeled by radiologists. Earlier, NIH Chest X-ray 14[30] proposed by Wang et al. contains 112,120 front-view images of 14 disease categories, among which there are 880 images of 8 categories containing box annotations. Lately, Nguyen et al. proposed VinDr-CXR [19] which contains 18,000 images that were manually annotated with 22 classes of rectangles surrounding abnormalities and 6 global labels of suspected diseases. There also exist some datasets that focus on a single disease, such as the Pneumonia detection dataset ¹¹1https://www.kaggle.com/c/rsna-pneumonia-detection-challenge, Tuberculosis detection dataset [18] and Pneumothorax segmentation dataset ²²2https://www.kaggle.com/c/siim-acr-pneumothorax-segmentation, etc. CAM[37], grad-CAM[26], and attention models[20, 35] are the most widely used tools to highlight rough locations of diseases. To use box annotations, Li et al. [14] propose an end-to-end fully convolutional neural network to classify and localize abnormalities. The problem is addressed in an approach like weakly supervised segmentation at the coarse level. Based on [14], Liu et al[16] propose a contrast induced attention network to predict the location of diseases, and an alignment module to align positive and negative images in the first place. Our work is different that we are the first to detect chest X-ray diseases with fully supervised annotations. Chen et al. [2] present a deep hierarchical multi-label classification approach for Chest X-ray diagnosis. Inspired by the idea of taxonomy in [2], we group 13 categories into three parent classes (Table 1) and evaluate the effect of each relation module based on the class hierarchy. More details are presented in the 4.6 subsection.

Modern object detection frameworks include two lines of approaches. The first line is two-stage detectors of Fast R-CNN[7], Faster R-CNN[23], FPN[15], etc. By using a top-down pathway and the lateral connections, FPN can enrich the semantic information of shallow layers while maintaining their high resolution. Utilizing multi-level features to improve performance is crucial in detection tasks. The second line is one-stage detectors of YOLO[22], SSD[17] and recent anchor-free detectors[13, 27, 39, 4, 38]. We choose FPN as our backbone framework since disease regions in chest X-rays range from 20$\times$ 20 pixels (e.g. nodule, calfication) to 1000$\times$ 1000 pixels (e.g. emphysema). FPN is known for extracting RoI features from shallow to deep layers, capturing characteristics from small to large objects. Mask R-CNN[9] is a successful instance segmentation framework extending FPN. The mask head operating on RoIs can output class-specific masks. Since disease regions are always of various shapes, we use mask R-CNN to visualize them for more accurate diagnosis. However, one difference for chest X-ray images is the constant anatomical structure, which is not available in natural images. Furthermore, general object detectors like FPN and Mask R-CNN fail to capture object-wise and class-wise relations, let alone object-structure relations. There exists works exploring relations for general object detection. In [11], appearance and geometry relations within objects are modeled inspired by attention modules[28] in NLP. In [33], object features are enhanced via attending different semantic concepts and propagating information through a common sense knowledge graph. In [32], a sparse graph is built based on semantic and spatial relations among objects. Our work is closely related to them while we focus on the relations between diseases and structures.

Some diseases or abnormalities are highly correlated with specific parts or organs in the body. For instance, cardiomegaly is a medical condition in which the heart is enlarged; atelectasis is a complete or partial collapse of the entire lung or area (lobe) of the lung.

Based on the above observation, we propose an anatomical structure relation module to encode spatial relations between disease and anatomical parts. To accomplish the task, we first adopt a pre-trained segmentation model ³³3https://github.com/Deepwise-AILab/ChestX-Det-Dataset/tree/main/pre-trained_PSPNet to obtain anatomical parts. Then we choose five key parts of left lung, right lung, left scapula, right scapula and heart, as shown in Figure 4. The segmentation model is trained on 1000 chest images labeled with 5 parts from external data, and is adopted to generate anatomical parts for each image in our datasets. Then for each disease RoI, we use coordinate difference to quantify its spatial relation with each anatomical part.

where $\left\{(x_{i}^{a},y_{i}^{a})\right\}_{i=1}^{2}$ and $\left\{(x_{i}^{p},y_{i}^{p})\right\}_{i=1}^{2}$ are up-left and bottom-right vertex coordinates of disease RoI and anatomical part respectively, and $w_{l}$ and $h_{l}$ are width and height of the bounding box covering left and right lung. Relating all five parts, $M=Concat\left\{(M_{p})_{p=1}^{5}\right\}$ is a 40-dimensional vector. Inspired by [28], we embed $M$ into a high-dimensional space, as

where sine and cosine functions of different wavelengths are computed. We can change the dimension of $f_{spa}$ by setting different $d_{e}$. Finally, the spatial relation between all disease proposals and anatomical parts are recorded in position embedding $f_{spa}(N_{r})\in\mathbb{R}^{N_{r}\times D_{spa}}$, where $D_{spa}=80*d_{e}$.

Contextual cues are very useful for radiologists to reference. For instance, by checking on the symmetric area on the contralateral side, radiologists can decide if the abnormality is a nodule or simply papilla. In our work, we focus on diseases in lung fields. Therefore we bound contextual area into the anatomical parts of left and right lung. Here we customize the attention mechanism having been proven successful in natural language processing and natural image recognition. Similarly, for a query of a particular disease RoI, a set of key contents in the lung fields are aggregated according to attention learned on disease-context relations. The relations involve both spatial and feature compatibility.

Specifically, we use $M_{l}$ and $M_{r}$ to denote the bounding box of left lung and right lung respectively. We adopt the same operation on $M_{l}$ and $M_{r}$ to obtain their grid-shaped feature $f_{l}$ and $f_{r}$, of size $N_{d}^{2}\times D_{M}$. For each disease RoI feature $f_{a_{i}}$ (after 2 fc layers) and each grid feature $f_{g_{j}}$, we compute their compatibility as

where $W_{a}$ and $W_{g}$ are both weight matrices, and $W_{a}f_{a}$ and $W_{g}f_{g}$ share the same dimension.

Similar with Eq.(1) and Eq.(2), we can obtain the spatial relation between abnormality $a_{i}$ and grid $g_{j}$.

where $s(a,g)$ is $(\frac{x_{1}^{a}-x^{g}}{w_{l}},\frac{y_{1}^{a}-y^{g}}{h_{l}},\frac{x_{2}^{a}-x^{g}}{w_{l}},\frac{y_{2}^{a}-y^{g}}{h_{l}})^{T}$ ⁴⁴4$(x_{i}^{a},y_{i}^{a})$ are coordinates of up-left and bottom-right vertexes of abnormality RoI. $(x^{g},y^{g})$ is the center coordinate of a grid., and $W_{s}\in\mathbb{R}^{1\times 4}$ is a learnable weight matrix.

Then, the spatial relation $\mathcal{P}(a_{i},g_{j})$ and appearance relation $\mathcal{F}(a_{i},g_{j})$ are collected across all grids and summed into a soft-max function to get attention weights as

Finally, grid-wise contextual features from lung fields are aggregated according to attention learned above.

where $W_{k}$ denotes the $1\times 1$ convolution kernel weights. The enhanced features of all region proposals are $f_{cxt}(N_{r})\in\mathbb{R}^{N_{r}\times D_{cxt}}$. Note that [,] means the concatenation operation.

Diseases or abnormalities in chest X-Rays are highly correlated. For instance, pulmonary tuberculosis is a complex disease which might contain nodules, fibrosis and consolidation simultaneously. There also exists causal relations among abnormalities. For instance, rib fracture is likely to cause pneumothorax. To this end, first we need to build a relation graph containing co-occurrence and causal relations among diseases. Then messages from semantic concepts of diseases are propagated via the relation graph. Finally RoI features aggregate messages to form representation $f_{cate}$. These three steps are detailed in the following.

Relation Graph Construction. To build the relation graph, we count co-occurrence frequency in the data set of each pair of categories and compute their conditional probability as

where $\#(B)$ denotes number of samples disease B exists, and $\#(A,B)$ denotes number of samples disease A and B co-exist. Then we build a directed graph ${G}=<N,E>$ as the disease relation graph. $N$ are disease category nodes and each edge $e_{i,j}\in E$ encodes relations:

Message Passing via Relation Graph. The relation graph is built based on statistics collected from thousands of X-rays and reflects prior knowledge. However, for a specific sample, it only reflects the condition of a patient in a limited period. Therefore, we propose to use a global attention to force message passing on potential diseases.

First, we input the whole X-ray image into the same convolutional layers of SAR-Net to get the image feature $F\in\mathbb{R}^{H\times W\times D^{{}^{\prime}}}$. Then a global mean-average-pooling operation is applied to squeeze $F$ into $F_{s}\in\mathbb{R}^{D^{{}^{\prime}}}$. Finally, we use a fully-connected layer to get categories score $\beta\in\mathbb{R}^{C}$, where $C$ is the number of categories. The global binary cross-entropy loss for each category is defined as

where $i$ is the index of categories, $y_{i}$ denotes the target label of the category and $p_{i}$ denotes the predicted probability.

Inspired by some few/zero-shot works [25, 29, 6], we use parameter weights of the classifier branch in box head to represent the semantic embedding of each category. Formally, the semantic embedding is defined as $W_{emb}\in\mathbb{R}^{C\times D}$, where $C$ and $D$ are number of categories and weight dimensions respectively. By transmitting causal or co-occurrence relations through the sparse relation graph, we can obtain the $i$-th category embedding as

Mapping Disease Relation to Regions. Since our goal is to enhance the original region features, we need to map the global embedded semantics from $C$ categories to $N_{r}$ regions. In our work, we choose the classification probability of each RoI $P\in\mathbb{R}^{N_{r}\times C}$ as the mapping bridge. The semantic embedding of $k$-th region is

where $p_{ki}$ denotes the probability for the $k$-th region towards category $i$.

Finally, $r_{k}$ is linearly transformed to reduce dimension

where $W_{t}$ is the weight of a fully-connected layer. The enhanced features of all regions are $f_{cate}(N_{r})\in\mathbb{R}^{N_{r}\times D_{cate}}$. Now, we obtain the embedded feature of spatial relation $f_{spa}$, contextual relation $f_{cxt}$ and category relation $f_{cate}$. The final enhanced feature of each disease RoI is $f^{\prime}=[f:f_{spa}:f_{cxt}:f_{cate}],f^{\prime}(N_{r})\in\mathbb{R}^{N_{r}\times(D+D_{spa}+D_{cxt}+D_{cate})}$.

Datasets: ChestX-Det is a subset with box annotations of NIH ChestX-14. NIH ChestX-14 contains 112,120 front-view images, among which there are 880 images with box annotations. To make fully supervised learning feasible, we select 3,575 images and invite three board-certified radiologists to annotate them with 13 common categories of diseases or abnormalities. For annotation, we split three board certified radiologists into two roles. Commitee: Every chest X-ray is annotated by the first two radiologists. The two radiologists are mutually blind to the annotation. Judge: The third and the most experienced radiologist with 15+ years of experience select annotations from those provided by the first two radiologists, and can add annotations if there are more. Figure 5 shows examples of the annotated 13 categories and a “normal” sample without abnormalities. We use 3025 images for training and 553 images for testing. 10% samples from the training set are used for validation.

DR-private is a private dataset by collecting 6,629 chest X-ray images from multiple Chinese hospitals. The annotation process is the same with ChestX-Det. As for DR-private, we use 5800 images for training and 829 images for testing. 10% samples from the training set are used for validating. Table 1 shows the instance count of each disease in both datasets. For both instance segmentation and anatomical part segmentation, all data are selected to cover the appearance range as wide as possible. We include normal and abnormal samples in different angles, whiteness, and scanning conditions. For white-out lungs in part segmentation, we ask radiologists to annotate their original contours as far as they can.

Part Segmentation Performance: We evaluate the part segmentation model on 229 test samples, the mean IOU (Mean Intersection-Over-Union) is 86.85%. We only use the bounding boxes of segmentation for further computation. The accuracy of segmentation does not seriously affect the bounding boxes. The segmentation results of PSP-Net can well serve the SAR-Net.

Evaluation: We borrow the metric of bounding box AP (AP^bb) [5] used in general object detection. Considering that disease/abnormality regions lack clear boundary as general objects, we adopt AP${}^{bb}_{50}$ as the main evaluation metric. AP${}^{bb}_{25}$ and AP${}^{bb}_{75}$ are used for reference. We believe that AP${}^{bb}_{75}$ better reflects localization performance and AP${}^{bb}_{25}$ better reflects classification performance. AP${}^{bb}_{50}$ is the comprehensive performance index for classification and localization. For more practical usage, we also present instance-level recall at fixed FP(false positive) per image for more direct reference. In addition, we use mask AP (AP^mask) to evaluate the performance of instance segmentation. AP${}^{mask}_{50}$ and AP${}^{mask}_{75}$ are also provided for reference.

We implement SAR-Net stacked on six backbones: ResNet-50-C4[10] in Faster R-CNN[23], Mask R-CNN[9] and Cascade R-CNN[1] respectively; ResNet-50-FPN[15] in Faster R-CNN and Mask R-CNN respectively; ResNet-50-FPN+DCN[3] in Mask R-CNN. The core idea of DCN is to capture position-based semantic context to improve the capacity of feature representation. It adds 2D offsets to the regular grid sampling locations and defines a new grid to select input values.

All experiments are implemented using Pytorch on 4 TITAN-V GPUs. ResNet-50 is pretrained on ImageNet [24]. For all training, we apply stochastic gradient descent (SGD) with a weight decay of 0.0001 and momentum of 0.9 to optimize all models. The first conv layers of FPN and C4 are frozen. We train 50 epochs with image batch-size of 2 on each GPU. The learning rate starts at 0.01, and reduce by a factor of 10 after 20 and 40 epochs.

During training, we adopt random flipping and multi-scale sampling (shorter side=$\left\{800\sim 1400\right\}$) for all images. At testing stage, the shorter side of image is fixed at 1200. The total number of proposed regions $N_{r}$ after NMS is 512. All the other hyper parameters and loss functions follow traditional Mask R-CNN. As shown in Table2, we clarify parameter dimensions in three modules with FPN-DCN. The extracted features $f_{cate}$, $f_{spa}$ and $f_{cxt}$ from three relation modules are all 1D. All the weight parameters are initialized using gaussian prior.

Table3 shows detection results comparison of our SAR-Net and baseline of various models on both datasets. We can see that SAR-Net consistently outperforms baseline with any configuration of models, evaluation metrics and datasets. Improvements on all models demonstrate that SAR-Net can be embedded into general two-stage object detection frameworks.

Stable gains are achieved at AP${}^{bb}_{50}$. We believe that SAR-Net plays an effective balance role in improving both classification and localization, which accords with the principle of SAR-Net. Using the strong model Mask R-CNN with FPN+DCN, SAR-Net can also boost baseline by 2.7% at AP${}^{bb}_{50}$ on ChestX-Det. It demonstrates that our modules are robust enough. Improvements on dataset DR-Private also demonstrates our effectiveness. It also means SAR-Net generalizes well on different data distributions. It is noteworthy that the results on DR-private are slightly lower than the results on ChestX-Det. The main reason leading to different performance on two datasets is different data distribution of the two datasets. ChestX-Det from ChestX-14 has more bedside samples and DR-private has more regular ones. Abnormalities in bedside samples are often severer and salient.

We also report the performance of instance segmentation on both datasets. As shown in Table 4, our relation modules can also achieve considerable performance gain in the instance segmentation task.

Table 5 shows comparison results at AP${}^{bb}_{50}$ of each category. From our observation, pleural thickening, fracture, diffusive nodule, cardiomegaly and emphysema are categories with consistent and large improvements on both datasets. For representative categories of fracture and diffusive nodule, our method achieves AP${}^{bb}_{50}$ of 38.9% and 45.6% on ChestX-Det, with a lead of 8.0% and 7.4% over baseline respectively. Fibrosis is the only category with decreasing accuracy in both datasets. In our datasets, emphysema and fibrosis often exist simultaneously and have similar spatial distributions, both of which diffuse the whole lung fields. Compared with fibrosis, the feature of emphysema is easier to distinguish, which cause that when fibrosis and emphysema overlap, the proposal boxes are easier to be predicted as emphysema. When adding the relation modules, such phenomenon is more obvious. The reason is that our SAR-Net has high sensitivity for the diseases which has more unique appearances, such as cardiomegaly and emphysema etc. Although the detection performance of fibrosis decreases slightly, the detection performance of emphysema obtains great improvement over both datasets. For more practical usage, we also provide recall (sensitivity) at fixed false positives per image for each category in Table 6. In particular, we set the rate of instance-level FP/image to be 0.1 at AP${}^{bb}_{25}$. From table 6, we can see that for most categories, our SAR-Net can obtain higher sensitivity than baseline.

We perform experiments with stronger backbones ResNet-101 and ResNeXt-101 [31]. Results are shown in Table 7.

Our proposed SAR-Net consistently improves performance on all the experimented backbones. Specially, on the ResNet-101 backbone, we can obtain the gains of 3.5, 2.9 and 3.2 points on the AP${}^{bb}_{25}$, AP${}^{bb}_{50}$ and AP${}^{bb}_{75}$ respectively. It demonstrates that our modules are robust enough and can be generalized to boost performance on multiple backbone architectures.

In order to evaluate the effectiveness of each module, we conduct ablation studies on ChestX-Det from different perspectives (4.5-4.7). In all experiments, we use the same train, validation and test set. ResNet-50 is adopted as the backbone. We remove the spatial relation module (SRM), the disease relation module (DRM) and the contextual relation module (CRM) from SAR-Net respectively. Ablation results are shown in Table 8.

The effect of spatial relation. Compared with SAR-Net, if we remove the spatial relation module, the performance is decreased by 0.0%, 0.4% and 1.5% on AP${}^{bb}_{25}$, AP${}^{bb}_{50}$ and AP${}^{bb}_{75}$ respectively. The maximum decline on AP${}^{bb}_{75}$ validates the effectiveness of the spatial relation module on localization. Moreover, the zero decline on AP${}^{bb}_{25}$ validates that the spatial relation module has little positive effect on classification, which corresponds to the principle of the spatial relation module.

The effect of disease relation. Compared with SAR-Net, removing the disease relation module decreases performance by 1.7%, 0.7% and 0.6% on AP${}^{bb}_{25}$, AP${}^{bb}_{50}$ and AP${}^{bb}_{75}$ respectively. This indicates that the disease relation module enhances feature representation more for classification than localization. When adding the disease relation module, some disease labels are corrected according to the information propagation among diseases.

The effect of contextual relation. Compared with the spatial relation+disease relation, adding the extra contextual relation module can boost the performance by 1.6%, 1.1% and 1.3% on AP${}^{bb}_{25}$, AP${}^{bb}_{50}$ and AP${}^{bb}_{75}$ respectively. The contextual relation module is encoded by spatial information and feature information. This encoding mechanism has effectiveness on both classification and localization, which is verified by the experimental results. Moreover, the proposed attention mechanism can obtain more gains.

In this subsection, we evaluate the performance of each relation module on different diseases. Specially, inspired by [2], we group all diseases into three super-classes, namely, LUNG, PLEURA and MEDIASTINUM. Detailed information of the super-classes and their sub-classes can be found in Table 1. Performance evaluations are conducted based on the super-classes and results are shown in Table 9.

The contextual relation module obtains maximum gains on AP${}^{bb}_{25}$, AP${}^{bb}_{50}$ and AP${}^{bb}_{75}$ respectively for LUNG diseases, while simply adding the spatial or disease relation module can not obtain great improvement. This is easy to understand, the diseases of LUNG distribute in the lung fields but most of them do not appear on fixed locations. For instance, the atelectasis may appears in either the upper lobe or in the lower lobe, given different causes. While the contextual relation module are able to models contextual relations between diseases and observations in the lung fields, which involves both spatial and appearance compatibility, and therefore is more effective for the diagnosis of lung diseases.

The diseases of PLEURA have strong co-occurrence relations. For instance, the effusion and primary spontaneous pneumothorax (PSP) often cause pleural thickening. Moreover, those three diseases have specific spatial distribution and locate around the pleura. Therefore, it is anticipated that the disease module and the spatial module should contribute more to the performance improvement. The experimental results also validated our hypothesis, as the disease module obtains gains on AP${}^{bb}_{25}$ & AP${}^{bb}_{50}$ and the spatial module obtain gains on AP${}^{bb}_{50}$ & AP${}^{bb}_{75}$. This suggests that the disease module can boost the performance for classification while the spatial module are better at improving localization accuracy, as we have emphasized before.

Compared with other diseases, cardiomegaly (MEDIASTINUM) has obvious appearance and fixed location, which makes it the perfect target disease of our proposed relation modules. It can be observed that all three relation modules can obtain considerable gains on AP${}^{bb}_{25}$ and AP${}^{bb}_{50}$. And the spatial relation module can also obtain gains on AP${}^{bb}_{75}$, a metric that requires precise localization.

Different relation modules have different emphasis (classification or localization) for different diseases. In general, when adding all the relation modules, our model can achieve maximum performance.

Model complexity. We also report the complexity of each module in Table 9. As shown in Table 9, using our modules can obtain higher performance with only a slight increase in parameters and floating point operations (FLOPs). Note that some parameters are shared in different modules, such as the coordinate parameters of the anatomical structure. Hence, the extra parameters of SAR-Net are less than the sum of extra parameters of all three modules. The same is true for FLOPs.

Figure 6 shows qualitative results of all the ablation methods. As shown in the first row, the disease relation module has little positive effect on localization. However, the spatial relation module and contextual module can effectively improve the localization performance of the detected diseases. This is because all the above two relation modules are better at exploiting the spatial information.

Pneumothorax and effusion have strong co-occurrence relations and specific spatial distribution and thus can benefit from the disease and spatial relation module. As shown in the second row, the disease relation module can correct the wrong disease labels according to the information propagation among diseases. On the other hand, the spatial relation module also can achieve the same effect with the help of the spatial information propagation. As have been shown on the quantitative results, the contextual relation module has high sensitivity on diagnosing the abnormalities of LUNG, but has low sensitivity on diagnosing the abnormalities of PLEURA. On this sample, the fracture is detected and pneumothorax is overlooked.

Pleural thickening develops when scar tissue thickens the delicate membrane lining the lungs (the pleura) and will not appear in the lung fields. As shown in the third row, the relation modules (SRM, CRM) which contain spatial information can eliminate this kind of false positives effectively. In addition, since the contextual relation module is not sensitive for PLEURA, it can not improve the localization performance of pleural thickening compared with the spatial relation module.

Moreover, we also show some failure cases in Figure 7. As mentioned before, effusion has strong co-occurrence relations with pneumothorax and often appears at costophrenic angle. As shown in the first case, when the relations modules are used, the above factors cause the consolidation box to be predicted as an effusion box. As mentioned in the subsection 4.3, when fibrosis and emphysema boxes overlap, the proposal boxes are easier to be predicted as emphysema. As shown in the second case, when the relations modules are used, the proposal box is predicted incorrectly and fibrosis is overlooked. In actual post-processing, this prediction box would be deleted by NMS. We keep it in Figure 7 for illustration purpose. There exist a lymphoid mass around the neck in the third case. In our dataset, mass often appears in the lung fields, which causes the lymphoid mass to be overlooked when the relation modules are present. Note that such failure cases only occur occasionally.

Extensive experiment results indicate that our proposed method is effective and has a great potential for clinical application. In Figure 2, we also show some instance segmentation results of SAR-Net on ChestX-Det for close inspection. Only the contours are shown to retain the appearance of the detected diseases.