Boundary Difference Over Union Loss For Medical Image Segmentation

The Boundary IoU is a segmentation evaluation metric which mainly focused on boundary quality. Given the ground truth binary mask $G$, the $G_{d}$ denotes the inner boundary region within the pixel width of $d$. The $P$ is the predicted binary mask, and $P_{d}$ denotes the corresponding inner boundary region, whose size is determined as a fixed fraction of 0.5% relative to the diagonal length of the image. Then, we can compute the Boundary IoU metric by using following equation, as shown in the left of Fig. 1,

A large Boundary IoU value indicates that the $G_{d}$ and $P_{d}$ are perfectly matched, which means $G$ and $P$ are with a similar shape and their boundary are well aligned. In practice, the $G_{d}$ and $P_{d}$ is computed by the erode operation [6]. However, the erode operation is non-differentiable, and we can not directly leverage the Boundary IoU as a loss function for training for increasing the consistency between two boundary areas.

As shown in the left of Fig. 1, we can find that the union of the two boundaries $|(G_{d}\cap G)\cup(P_{d}\cap P)|$ actually are highly correlated to the difference set between $G$ and $P$. The intersection $|(G_{d}\cap G)\cap(P_{d}\cap P)|$ is correlated to the inner boundary of the intersection of $G$ and $P$. If the difference set $G\cup P-G\cap P$ decreases and the $G\cap P$ increases, the corresponding Boundary IoU will increase.

Based on the above analysis, we design a Boundary DoU loss based on the difference region to facilitate computation and backpropagation. First, we directly treat the difference set as the miss-matched boundary between $G$ and $P$. Besides, we consider removing the middle part of the intersection area as the inner boundary, which is computed by $\alpha*G\cap P~{}(\alpha<1)$ for simplicity. Then, we joint compute the $G\cup P-\alpha*G\cap P$ as the partial union. Finally, as shown in the right of Fig. 1, our Boundary DoU Loss can be computed by,

where $\alpha$ is a hyper-parameter controlling the influence of the partial union area.

Adaptive adjusting $\alpha$ based-on target size: On the other aspect, the proportion of the boundary area relative to the whole target varies for different sizes. When the target is large, the boundary area only accounts for a small proportion, and the internal regions can be easily segmented, so we are encouraged to focus more on the boundary area. In such a case, using a large $\alpha$ is preferred. However, when the target is small, neither the interior nor the boundary areas are easily distinguishable, so we need to focus simultaneously on the interior and boundary, and a small $\alpha$ is preferred. To achieve this goal, we future adaptively compute $\alpha$ based on the proportion,

where $C$ denotes the boundary length of the target, and $S$ denotes its size.

In this part, we compare Boundary DoU Loss with Dice Loss. Firstly, we can re-write our Boundary DoU Loss as:

where $S_{D}$ denotes the area of the difference set between ground truth and prediction, $S_{I}$ denotes the intersection area of them, and $\alpha^{\prime}=1-\alpha$. Meanwhile, the Dice Loss can be expressed by the following:

where TP, FP and FN denote True Positive, False Positive, and False Negative, respectively. It can be seen that Boundary DoU Loss and Dice loss differ only in the proportion of the intersection area. Dice is concerned with the whole intersection area, while Boundary DoU Loss is concerned with the boundary since $\alpha<1$. Similar to the Dice loss function, minimizing the $L_{DoU}$ will encourage an increase of the intersection area ($S_{I}\uparrow$) and a decrease of the different set ($S_{D}\downarrow$). Meanwhile, the $L_{DoU}$ will penalize more over the ratio of ${S_{D}}/{S_{I}}$. To corroborate its effectiveness more clearly, we compare the values of $L_{Dice}$ and $L_{DoU}$ in different cases in Fig. 2. The $L_{Dice}$ decreases linearly with the difference set, whereas $L_{DoU}$ will decrease faster when $S_{I}$ is higher enough.

Synapse:¹¹1https://www.synapse.org/#!Synapse:syn3193805/wiki/217789 The Synapse dataset contains 30 abdominal 3D CT scans from the MICCAI 2015 Multi-Atlas Abdomen Labeling Challenge. Each CT volume contains $86\sim 198$ slices of $512\times 512$ pixels. The slice thicknesses range from 2.5 mm to 5.0 mm, and in-plane resolutions vary from $0.54\times 0.54~{}mm^{2}$ to $0.98\times 0.98~{}mm^{2}$. Following the settings in TransUNet [3], we randomly select 18 scans for training and the remaining 12 cases for testing.

We conduct experiments on three advanced models to evaluate the performance of our proposed Boundary DoU Loss, i.e., UNet, TransUNet, and Swin-UNet. The models are implemented with the PyTorch toolbox and run on an NVIDIA GTX A4000 GPU. The input resolution is set as $224\times 224$ for both datasets. For the Swin-UNet [2] and TransUNet [3], we used the same training and testing parameters provided by the source code, i.e., the learning rate is set to 0.01, with a weight decay of 0.0001. The batch size is 24, and the optimizer uses SGD with a momentum of 0.9. For the UNet, we choose ResNet50 as the backbone and initialize the encoder with the ImageNet pre-trained weights following the setting in TransUNet [3]. The other configurations are the same as TransUNet. We train all models by 150 epochs on both Synapse and ACDC datasets.

We further train the three models by different loss functions for comparison, including Dice Loss, Cross-Entropy Loss (CE), Dice+CE, Tversky Loss [19], and Boundary Loss [13]. The training settings of different loss functions are as follows. For the $\lambda_{1}\mbox{Dice}+\lambda_{2}\mbox{CE}$, we set $(\lambda_{1},\lambda_{2})$ as (0.5, 0.5) for the UNet and TransUNet, and (0.6, 0.4) for Swin-UNet. For the Tversky Loss, we set $\alpha=0.7$ and $\beta=0.3$ by referring to the best performance in [19]. Following the Boundary Loss [13], we use $L=\alpha*(\mbox{Dice}+\mbox{CE})+(1-\alpha)*\mbox{Boundary}$ for training . The $\alpha$ is initially set to 1 and decreases by 0.01 at each epoch until it equals 0.01.

Table 2 reports the results on the ACDC dataset. We can find that our Boundary DOU Loss effectively improves DSC on all three models. Compared with Dice Loss, the DSC of UNet, TransUNet, and Swin-UNet improved by 0.62%, 0.6%, and 0.85%, respectively. Although our Loss did not get all optimal performance for the Hausdorff Distance, we substantially outperformed all other losses on the Boundary IoU. These results indicate that our method can better segment the boundary regions. This capability can assist doctors in better identifying challenging object boundaries in clinical settings.