A CT-Based Airway Segmentation Using U²-net Trained by the Dice Loss Function

The ATM’22 provides 499 three dimensional CT scans from multi-sites, of which 299 are given for training, 50 for validation and the rest 150 will be used by the challenge organizer for testing the submitted algorithms. The axial slices of the provided CT scans are of the size $512\times 512$ in length and width but with varying depth sizes.

Given a chest training CT image, we first normalize the voxel values to prepare for faster computations in the following network training. The original distribution of the CT voxel intensity is around $[-1024,3071]$. We employ a CT-dependent normalization, that is, the intensity value is normalized according to the mean and standard deviation values of each slice. The normalized voxel intensity is rescaled to around $[0,2.5]$. Note that the irrelevant dark parts with low intensities in a CT image, which normally enjoy the constant standard deviation value, are further reset to a different value, for example, $-9$ to distinguish from the rescaled array value $0$.

In our training, we adopt the original size $512\times 512$ of each slice with the GPU memory permitting. However, the depth sizes are doubled, that is, each input slice image is of $512\times 512\times 2$ to improve the performance accuracy of the employed U²-net in the feature extraction. Although the treatment may demand more GPU memory, we have shown, with the comparison to the input images in the original depth size, that the double size treatment indeed significantly increases the training accuracy.

In the training process we employ the Uⁿ-net architecture [1] to obtain the saliency probability function for the airway tree. The Uⁿ-net is a n-level nested U-structure, where the exponent $n$ denotes the the level of the nested U-structures. We adopt $n=2$ in the application from the practical perspective by following [1]. The U²-net is built by two levels of U-structures. In details, the top level U-structure consists of a six-stage encoder, a five-stage decoder, and a saliency map fusion module and each stage among the encoders and decoders contains the newly proposed Residual U-block (abbreviated as RSU), which is the nested bottom level downsampling upsampling encoder-decoder U-net. Each level of the downsampling/upsampling paths in the corresponding RSU consists of a $3\times 3$ convolutional layer, a batch normalization layer, a rectified linear unit (ReLu) activation function, and a $2\times 2$ max-pooling/upsampling layer. The nested RSU allows each block to use the multiscale features as residuals instead of the original features, which makes the whole architecture maintain high resolution feature with a reasonable requirement of the GPU memory and computing ability.

While the input data are processed by each RSU, the corresponding encoder and decoder with a $3\times 3$ convolution layer, followed by a sigmoid function in each layer of the upper level U-net will generate six side output saliency probability maps $S_{side}^{(6)},S_{side}^{(5)},S_{side}^{(4)},S_{side}^{(3)},S_{side}^{(2)},S_{side}^{(1)}$, by applying a $1\times 1$ convolution layer and a sigmoid function to the concatenation of which leads to a final saliency probability map $S_{fuse}$. In addition to the economical cost of GPU memory and computing capacity, the training of the U²-net does not rely on any backbones due to the aforementioned nested design.

As mentioned in Subsection 2.1, the input is $512\times 512\times 2$ duplicated slices, hence, we change the channel number of the side probability output at each level to 3 to improve the performance of the architecture in extracting information and features at different levels and facilitate the recognition of the trachea as well as the small peripheral bronchi.

The training of the U²-net follows an end-to-end supervised manner, whose training loss is defined as the weighted sum of the loss of the side and final output saliency probability maps

where $l_{side}^{(n)}$ ($n=6$, since there are 6 side output in total) is the loss of the side output probability map, $l_{fuse}$ is the loss of the final output, and $w_{side}$ as well as ${w_{fuse}}$ are the associated weights of each loss. More specifically, we utilize the Dice loss function for each $l$ term (including both $l_{side}$ and $l_{fuse}$),

where $M$ denotes the total voxel index in a ground truth, $P_{G(i)}$ is the numerical array value of the ground truth and $P_{S(i)}$ represents the corresponding predicted value of the saliency probability map, $\epsilon$ is a coefficient used to avoid division by zero and we take $\epsilon=1$. The network is then trained via optimizing the overall loss $L$ by evaluating the network output with the given ground truth.

We implement a threshold of 0.5 on the output map $S_{fuse}$ to obtain the predicted mask for the 50 CT for validation. However, the predicted mask may contain spurious non-airway regions. Notably, the airway and non-airway regions can be generally topologically viewed in a Hausdorff space (see Fig. 2 for example). We then refine the segmented results by removing the distinguishable large non-airway parts. Due to the separability, we assign different labels to each separate region, for example, $1$ for the airway region and $2$ for the non-airway region as shown in Fig. 3. Similar post-processing methods can be found in [6]. Extracting the mask numerically labelled with $1$ and applying to the validation CT scans, we finally obtained the refined airway tree segmentation results.