Semi-Supervised Segmentation of Radiation-Induced Pulmonary Fibrosis from Lung CT Scans with Multi-Scale Guided Dense Attention

1). 2.5D Network Structure:

To deal with the anisotropic 3D resolution with high inter-slice spacing and low intra-slice spacing, we combine 2D (i.e., intra-slice) and 3D convolutions so that the network has an approximately balanced physical receptive field along each axis. Let S represent the number of scales in the network (S = 5 in this paper), we use 2D convolutions and 2D max-poolings for the first M scales in the encoder, and employ 3D convolutions and 3D max-poolings for the last S – M scales in the encoder. Each resolution level in the decoder contains the same type of 2D or 3D convolutional blocks as in the encoder. We use trilinear interpolation for upsampling in the decoder.

As our lung CT images have a resolution around 0.3×0.3×1.25 mm, i.e., the in-plane resolution is about four times of the through-plane resolution, we set M = 2 so that the 2D max-pooling layers in the first two scales make the resulted feature maps have a near-isotropic 3D resolution, as shown in Fig. 2. Note that the first and S-th scales have the highest and lowest spatial resolution, respectively. We use $F_s^e$ and $F_s^d$ to denote the feature maps obtained by the last convolutional block at scale s in the encoder and decoder, respectively. Note that $F_S^e=F_S^d$ in the bottleneck block.

2). Multi-Scale Guided Dense Attention:

To better deal with PF lesions with various position and size, we use multi-scale attention to improve the network’s spatial awareness, and propose dense attention to leverage multi-scale contextual information for the segmentation task. Specifically, at each scale s of the decoder, we use a convolutional layer to get a spatial attention map ${\tilde{P}}_s$, and use ${\tilde{P}}_s$ as a high-level attention signal to guide the learning in all lower scales of the decoder. The input of the decoder at scale s is a concatenation of three parts: $F_s^e$ from scale s of the encoder, an upsampled version of the decoder feature map $F_{s+1}^d$ and ${\tilde{P}}_{s+1}^s\oplus {\tilde{P}}_{s+2}^s\oplus \dots \oplus {\tilde{P}}_S^s$, where ⨁ is the concatenation operation and ${\tilde{P}}_{s+1}^s$ is the upsampled version of ${\tilde{P}}_{s+1}$ so that it has the same spatial resolution as $F_s^e$. Therefore, a lower scale accepts the attention maps from all the higher scales as input, which is referred to as Multi-Scale Dense Attention (MSDA). The decoder thus takes advantage of multi-scale contextual information that enables the network to pay more attention to the target region.

To better learn the spatial attention, we propose Multi-Scale Guided Attention (MSGA) to explicitly supervise ${\tilde{P}}_1,{\tilde{P}}_2,\dots ,{\tilde{P}}_S$ at different scales. Let P_s denote the softmaxed output of ${\tilde{P}}_s$ and Y denote the ground truth (i.e., one-hot probability map) of a training sample X. Unlike a common deep supervison strategy that upsamples ${\tilde{P}}_s$ or P_s at different scales to the same spatial resolution as Y [40]–[42], we down-sample Y to obtain multi-scale ground truth for spatial attention, which makes the loss calculation more efficient and can directly supervise the spatial attention maps at different scales. Let Y_s denote the down-sampled ground truth at scale s. The multi-scale loss function for a single image is:

where $\mathcal{P}=\left\{P_1,P_2,\dots ,P_S\right\}$. L_s() is a base loss function for image segmentation, such as the Dice loss [19]. α_s is the weight of L_s() at scale s.

1). Inference for Unannotated Images with Uncertainty Estimation:

With the network parameter θ_k−1 obtained in the last round, in round k, we first use θ_k−1 to predict provisional pseudo labels for unannotated images in ${\mathcal{D}}^U$ and the associated uncertainty estimation. Note that in the first round (i.e., k = 1), the initial network parameter θ₀ is obtained by pre-training with the annotated images in ${\mathcal{D}}^L$.

With θ_k−1, we employ the Monte Carlo (MC) Dropout [43] that has shown to be an effective method for estimation of epistemic uncertainty caused by the lack of training data [34], [36]. MC Dropout feeds an image $X_i^U$ into the network R times with random dropout, which leads to R predictions, i.e., R foreground probability maps [43]. The average of these R foreground probability maps is taken as the provisional probability map $P_i^U$, a binarization of which leads to a provisional pseudo segmentation label $Y_i^U$.

At the same time, the statistical variance of the R foreground probability maps is taken as the uncertainty map $V_i^U$, which gives voxel-level uncertainty. As uncertainty information can indicate potentially wrong segmentation results [34], [44], we treat the pseudo labels with high uncertainty (i.e., low confidence) values as unreliable labels for images in ${\mathcal{D}}^U$, and propose a Confidence-Aware Refinement (CAR) method to improve the pseudo labels’ quality.

2). Confidence-Aware Refinement of Pseudo Labels:

Given the provisional pseudo label $Y_i^U$ with the uncertainty map $V_i^U$ for an unannotated image $X_i^U\in {\mathcal{D}}^U$, and let x denote a voxel, we split the voxels in $Y_i^U$ into three sets according to the status: high-confidence foreground voxels $\mathbb{F}=\left\{\mathbf{x}\mid Y_{ix}^U=1,V_{i\mathbf{x}}^U\le ϵ\right\}$, high-confidence background voxels $\mathbb{B}=\left\{\mathbf{x}\mid Y_{ix}^U=0,V_{ix}^U\le ϵ\right\}$ and undetermined voxels $\mathbb{U}=\left\{\mathbf{x}\mid V_{ix}^U>ϵ\right\}$, where ϵ = 0.05 is a small threshold value. For undetermined voxels in $\mathbb{U}$, we refine their labels according to a contextual regularization considering intervoxel connections and softened probabilities for these voxels.

Our CAR for pseudo label refinement has two steps: probability map softening and contextual regularization. First, we soften the foreground probabilities for uncertain voxels, which will degrade the influence of the network’s prediction for these voxels in the following contextual regularization step. Let p denote a foreground probability value and u denote the corresponding uncertainty value, the softening function is:

where the softened foreground probability get closer to 0.5 when u is larger. Let ${\dot{P}}_{ix}^U$ denote the softened foreground probability of voxel x, and is obtained by:

Then, we use contextual regularization taking ${\dot{P}}_i^U$ as input to refine the pseudo labels, which is implemented by a fully connected Conditional Random Field (CRF) [45]. For simplicity, we denote $X_{i\mathbf{x}}^U$, $Y_{i\mathbf{x}}^U$ and ${\dot{P}}_{ix}^U$ as x_x, y_x and ${\dot{p}}_{\mathbf{x}}$, respectively. The energy function of CRF is:

where $\varphi \left(y_{\mathbf{x}}\right)=-y_{\mathbf{x}}\log \left({\dot{p}}_{\mathbf{x}}\right)-\left(1-y_{\mathbf{x}}\right)\left(1-{\dot{p}}_{\mathbf{x}}\right)$ constrains the output to be consistent with the softened probability map, and this constraint for the low-confidence voxels in $\mathbb{U}$ is weak. The second term in Eq. (4) is a pairwise potential that encourages the label’s contextual consistency:

where μ(y_x, y_y) = 1 if y_x ≠ y_y and 0 otherwise. Minimization of Eq. (4) leads to a refined pseudo label $Y_i^U$ for $X_i^U$.

3). Confidence-based Image Weighting of Pseudo Labels:

With the new pseudo label $Y_i^U$, we further employ the confidence to update its image-level weight $w_i^U$ to suppress low-quality pseudo labels at the image level. We first define an image-level uncertainty v_i based on uncertainty map $V_i^U$:

where v_i is the sum of voxel-level uncertainty normalized by the segmented lesion’s volume in the image. η = 10⁻⁵ is a small number for numerical stability. Let v_max and v_min denote the maximal and minimal values of v_i among all the unannotated images, we map v_i to the range of [0, 1]:

Finally, the image-level weight $w_i^U$ is defined as:

where γ ≥ 1.0 is a hyper parameter to control the non-linear mapping between the image-level uncertainty $v_i^{\prime }$ and the weight. We do not use γ < 1.0 as it will lead the weight of most samples to be very small (close to 0.0). In contrast, γ > 1.0 leads the weight for most samples close to 1.0, and only samples with a high uncertainty will be strongly suppressed. In the experiment, we set γ = 3.0 according to the best performance on the validation set.

4). Model Update with Batch Training:

With the refined pseudo labels and image-level weights obtained above, we train the network based on ${\mathcal{D}}^L$ and the current pseudo labels for images in ${\mathcal{D}}^U$, where each image is weighted by $w_i^L$ or $w_i^U$ in the segmentation loss function. The weighted loss for the entire training set is:

where ℓ() is defined in Eq. (1). ${\mathcal{P}}_i^L$ and ${\mathcal{P}}_i^U$ are the multi-scale predictions obtained by PF-Net for an annotated image and an unannotated image, respectively.

1). Data:

Thoracic CT scans of 41 Male Rhesus Macaques with radiation-induced lung damage were collected with ethical committee approval. Once irradiated, each individual underwent serial CT scans for assessment of PF around every 30 days in 3 to 8 months. 133 CT scans with PF were used for experiments of the segmentation task. The CT scans have a slice thickness of 1.25 mm, with image size 512 × 512 and pixel size ranging from 0.20 mm × 0.20 mm to 0.38 mm × 0.38 mm. We randomly split the dataset at individual level into 86 scans from 27 individuals for training, 15 scans from 4 individuals for validation, and 32 scans from 10 individuals for testing. Manual annotations given by an experienced radiologist were used as the segmentation ground truth. In the training set, we used 18 scans with annotations as ${\mathcal{D}}^L$ and the other 68 scans as unannotated images ${\mathcal{D}}^U$ for semi-supervised learning, and also investigated the performance of our method with different ratios of annotated images. For preprocessing, we crop the lung region and normalize the intensity to [0, 1] using a window/level of 1500/−650.

2). Implementation and Evaluation Metrics:

Our PF-Net¹ and I-CRAWL framework were implemented in Pytorch with PyMIC² [14] library on a Ubuntu desktop with an NVIDIA GTX 1080 Ti GPU. The channel number parameter N in PF-Net was set as 16. Dropout was only used in the encoder. The dropout rate at the first two scales of the encoder was 0 due to the small number of feature channels, and that for the last three scales was 0.3, 0.4 and 0.5, respectively. We set the base loss function L_s() as Dice loss [19]. PF-Net was trained with Adam optimizer, weight decay of 10⁻⁵, patch size of 48 × 192 × 192 and batch size of 2.

The round number for our I-CRAWL was K = 3, and we use round 0 to refer to the pre-training with annotated images. In each of the following round, the pseudo labels of unannotated images acquired by CAR were kept fixed, and we used Adam optimizer to train the network for tens of thousands of iterations until the performance on validation set stopped to increase. The learning rate was initialized as 10⁻³ and halved every 10k iterations. Uncertainty estimation was based on 10 forward passes of MC dropout. We used the SimpleCRF library³ to implement fully connected CRF [45]. Following [46], image intensity was rescaled from [0, 1] to [0, 255] before the image was sent into the CRF, and the CRF parameters were: w₁ = 3, w₂ = 10, σ_α = 10, σ_β = 20 and σ_γ = 15, which was tuned based on the validation set. γ in Eq. (8) was 3.0, and the performance based on different γ values is shown in Table III.

For quantitative evaluation of the segmentation, we used Dice score, Relative Volume Error (RVE) and the 95-th percentile of Hausdorff Distance (HD₉₅) between the segmented PF lesions and the ground truth in 3D volumes. Paired t-test was used to see if two methods were significantly different.

1). Comparison with Existing Networks:

Our PF-Net was compared with three different categories of network structures: 1) 2D CNNs including the typical 2D U-Net [15] and more advanced attention U-Net [22] that leverages spatial attention to focus more on the segmentation target; 2) 3D CNNs including 3D U-Net [18], 3D V-Net [19] and scSE-Net [17] that combines a 3D U-Net [18] backbone with spatial and channel “squeeze and excitation” modules; and 3) existing networks designed for dealing with volumetric images with large inter-slice spacing (i.e., anisotropic resolution): nnU-Net [27] that automatically configures the network structure so that it is adapted to the given dataset, AH-Net [25] that transfers features learned from 2D images to 3D anisotropic volumes, and VS-Net [29] that uses a mixture of 2D and 3D convolutions with spatial attention. Quantitative evaluation results of these methods are shown in Table I. We found that 3D U-Net [18] achieved a Dice score of 57.19%, which was the lowest among the compared methods. 2D U-Net [15], Attention U-Net [22], 3D V-Net [19] and 3D scSE-Net [17] achieved similar performance with Dice score around 60%. The nnU-Net [27], AH-Net [25] and VS-Net [29] designed to deal with anisotropic resolution generally performed better than these 2D and 3D networks. Among these existing methods, VS-Net achieved the best Dice that was 67.45%. Our PF-Net achieved Dice, RVE and HD₉₅ of 70.36%, 27.96% and 10.87 mm, respectively, where the Dice and HD₉₅ were significantly better than those of the other compared methods.

Fig. 4 shows a qualitative comparison between these networks, where (a) and (b) are from two individuals, and axial and coronal views are shown in each case. In Fig. 4(a), it can be seen that 2D U-Net [15], 3D V-Net [19] and AH-Net [25] lead to obvious under-segmentation, as highlighted by blue arrows in the first row. VS-Net [29] performs better than them, but the result of our PF-Net is closer to the ground truth than that of VS-Net. From the coronal view of Fig. 4(b), we can observe that the result of 2D U-Net [15] lacks inter-slice consistency. 3D V-Net [19] achieves better inter-slice consistency, but it has a poor segmentation in the upper and lower regions of the lesion, as indicated by the blue arrow in the third column. AH-Net [25], VS-Net [29] and our PF-Net that consider anisotropic resolution obtain better performance than the above networks purely using 2D or 3D convolutions. What’s more, Fig. 4 shows that the lesions have complex and irregular sizes and shapes, and the proposed PF-Net is able to achieve more accurate segmentation in these cases than AH-Net [25] and VS-Net [29].

2). Ablation Study:

To investigate the effectiveness of each component of our PF-Net, we set the baseline as a naive 2.5D U-Net that extends 2D U-Net [15] by using pReLU, replacing 2D convolutions with 3D convolutions at the three lowest resolution levels, and adding dropout layers to each convolutional block. To justify the choice of using 2D convolutions at the first two resolution levels and 3D convolutions at the other resolution levels, we set M to 0–5 respectively. Note that M = 0 and M = 5 correspond to pure 2D and pure 3D networks, respectively. Comparison between these variants are listed in the first section of Table II, which shows that the performance increases as M changes from 0 to 2, and decreases when M is 3 and larger. This is in line with our motivation to set M to 2 due to the fact that the in-plane resolution is around four times of the through-plane resolution. Thus, we use baseline (M = 2) in the following ablation study.

The proposed PF-Net is referred to as baseline + MSDA + MSGA, where MSDA is our proposed multi-scale dense attention and MSGA is the proposed multi-scale guided attention. We compared the baseline and our PF-Net with: 1) Baseline + non-local [23], where the non-local is a self-attention block inserted at the bottleneck (scale 5) of the baseline network, and it was not used at the lower scales with higher resolution due to memory constraint; 2) Baseline + deep supervision, where the baseline network was only combined with a typical deep supervision strategy as implemented in [40], [41]; 3) Baseline + MSGA, without using MSDA; 4) Baseline + MSGA + MSDA°, where MSDA° is a variant of MSDA and it refers to ${\tilde{P}}_s$ is only sent to its next lower scale s − 1, rather than all the lower scales in the decoder of PF-Net.

Quantitative evaluation results of these variants are listed in the second section of Table II. It shows that compared with the baseline (M = 2), deep supervision improved the Dice score from 67.87% to 68.57%, but our MSGA was more effective with Dice of 69.32%. Using MSDA° could further improve the performance, but it was less effective than our MSDA. Table II demonstrates that our PF-Net was better than the compared variants, and it significantly outperformed the baseline in terms of Dice and HD₉₅. Fig. 6 presents a visualization of multi-scale attention maps of PF-Net. It shows that attention maps across scales are consistent with each other and they change from coarse to fine as the spatial resolution increases.

1). Hyper-Parameter Setting:

First, to investigate the best value of hyper parameter γ in Eq. (8) that controls the weight of unannotated images, we measured the model’s performance on the validation set at the first round of I-CRAWL with γ ranging from 1.0 to 5.0. They are compared with “Initial” that refers to model pre-trained on the annotated images and “No weighting” that denotes treating all unannotated images equally without considering the quality of pseudo labels. Quantitative measurements shown in Table III demonstrate that the best γ value was 3.0, and its corresponding Dice score was 68.75%, which was better than 66.62% obtained by “Initial” and 67.83% obtained by “No weighting”. Therefore, we set γ = 3.0 in the following experiments.

2). Uncertainty and Confidence-Aware Refinement:

Fig. 5 shows a visual comparison between a standard fully connected CRF [45] and our proposed CAR that leverages confidence (uncertainty) for refinement of pseudo labels. The first row of each subfigure shows an unannotated image and the pseudo label with uncertainty obtained by the CNN, and the second row shows the updated pseudo labels. In Fig. 5(a), the initial pseudo label has a large under-segmentated region, which is associated with high values in the uncertainty map, i.e., low confidence. The standard CRF only fixed the pseudo label moderately. In contrast, with the help of confidence, our CAR largely improved the pseudo label’s quality by recovering the under-segmented region. In Fig. 5(b), the initial pseudo label has some over-segmentation in airways, and the uncertainty map indicates the potentially wrong segmentation in the corresponding regions well. We can observe that CAR outperformed CRF in removing the over-segmented airways.

We also compared MC Dropout [43] with two other uncertainty estimation methods: entropy minimization and Bayesian network [47]. A visual comparison of them is shown in Fig. 7. We found that in spite of longer time required than the other methods, MC Dropout could generate more calibrated uncertainty estimation. As shown in Fig. 7, the uncertain region obtained by entropy minimization and Bayesian network [47] are mostly located around the border of the segmentation output, and uncertainty map obtained by MC dropout can better indicate under- and over-segmentation regions, which is highlighted by red and yellow arrows in Fig. 7. For quantitative comparison, we applied CAR as a post processing method to the validation set in the first round of I-CRAWL, where these uncertainty estimation methods were used respectively. Results in Table IV show that using MC dropout for CAR improved the prediction accuracy from 66.62% to 69.05% in terms of Dice, which outperformed using the other two uncertainty estimation methods.

Fig. 8 shows a visualization of pseudo labels and uncertainty as the round increases. It can be observed that the initial pseudo label at round 0 has a large under-segmented region with high uncertainty. The pseudo label becomes more accurate and confident as the round increases.

3). Ablation Study of I-CRAWL:

For ablation study of I-CRAWL, we set the PF-Net trained only with the annotated images as a baseline, and it was compared with: 1) IT that refers to naive iterative training, where in each round pseudo label of an unannotated image is reset to the prediction given by the network without refinement, 2) IT + CRF that uses standard fully connected CRF [45] to refine pseudo labels, 3) IT + CAR denoting that our confidence-aware refinement is used to update pseudo labels in each round, and 4) our I-CRAWL (IT + CAR + IW) where IW denotes our confidence-based image weighting of pseudo labels.

The performance of these methods at different rounds are shown in Fig. 9. Note that round 0 is the baseline, and all the methods based on iterative training performed better than the baseline. However, the improvement obtained by only using IT is slight. Using CRF or CAR to refine the pseudo labels at different rounds achieved a large improvement of Dice, and our CAR considering the voxel-level confidence of predictions outperformed the naive CRF. Weighting of pseudo labels based on image-level confidence helped to obtain more accurate result, and our I-CRAWL outperformed the other variants. Fig. 9 also shows that the improvement from round 0 to round 1 of I-CRAWL is large, but the model’s performance does not change much at round 2 and 3.

Table V shows quantitative comparison between the baseline and variants of I-CRAWL at the end of training (round 3). It can be observed that IT’s performance was not far from the baseline, with an average Dice of 70.87% compared with 70.36%. Using CRF and CAR improved the average Dice to 72.13% and 72.71%, respectively, showing the superiority of CAR compared with CRF. I-CRAWL improved the average Dice to 73.04%, with HD₉₅ value of 7.92 mm in average, which was better than the other variants.

4). Comparison with Existing Methods:

I-CRAWL was compared with several sate-of-the-art semi-supervised methods for medical image segmentation: 1) Fan et al. [13] that uses a randomly selected propagation strategy for semi-supervised COVID-19 lung infection segmentation, 2) Bai et al. [32] that uses CRF to refine pseudo labels in an iterative training framework, corresponding to “IT + CRF” described previously, 3) Cui et al. [33] that is an adapted mean teacher method, and 4) UA-MT [34] that is uncertainty-aware mean teacher. For all these methods, we used our PF-Net as the backbone network.

We investigated the performance of these methods with different ratios of labeled data: 10%, 20% and 50%. For each setting, the baseline was learning only from the labeled images, and the upper bound was “full supervision” where 100% training images were labeled. The results are shown in Table VI. It can be observed that with only 10% images annotated, the baseline performed poorly with an average Dice of 54.73%. Our I-CRAWL improved it to 67.18%, which largely outperformed the other methods. When 20% images were annotated, our method improved the average Dice from 70.36% to 73.04% compared with the baseline, while the mean teacher-based methods did not bring much performance gain. When 50% images were annotated, our method also outperformed the others, and it was comparable with full supervision (74.02±13.77% versus 74.58±10.99% in terms of Dice, with p-value > 0.05).