Boundary Loss for Remote Sensing Imagery Semantic Segmentation

TL;DR

This paper introduces a new boundary-aware loss function for remote sensing image segmentation, improving boundary accuracy and overall segmentation quality in neural network models.

Contribution

The paper proposes a novel differentiable boundary loss function that enhances boundary delineation in remote sensing image segmentation tasks.

Findings

Models trained with the new loss outperform baselines in IoU scores.

The loss improves boundary accuracy in synthetic and real-world datasets.

Enhanced boundary detection leads to better segmentation quality.

Abstract

In response to the growing importance of geospatial data, its analysis including semantic segmentation becomes an increasingly popular task in computer vision today. Convolutional neural networks are powerful visual models that yield hierarchies of features and practitioners widely use them to process remote sensing data. When performing remote sensing image segmentation, multiple instances of one class with precisely defined boundaries are often the case, and it is crucial to extract those boundaries accurately. The accuracy of segments boundaries delineation influences the quality of the whole segmented areas explicitly. However, widely-used segmentation loss functions such as BCE, IoU loss or Dice loss do not penalize misalignment of boundaries sufficiently. In this paper, we propose a novel loss function, namely a differentiable surrogate of a metric accounting accuracy of boundary detection. We can use the loss function with any neural network for binary segmentation. We performed validation of our loss function with various modifications of UNet on a synthetic dataset, as well as using real-world data (ISPRS Potsdam, INRIA AIL). Trained with the proposed loss function, models outperform baseline methods in terms of IoU score.

Tables

Table 1. Table 1: I ​ o ​ U 𝐼 𝑜 𝑈 IoU values for different datasets and loss functions. ConvNN is a not very deep neural network with 5 convolutional blocks in both the encoder and the decoder. For ISPRS dataset UNet model with ResNet34 backbone was used, for INRIA dataset – UNet with Inception-ResNet-v2 backbone

	AICD	ISPRS Potsdam	INRIA AIL
	ConvNN	ResNet34	Inc.-ResNet-v2
$L_{B{F}_1,IoU}$	86.91	93.85	74.30
$L_{IoU}$	85.66	92.56	71.91
$L_{SS,IoU}$	${61.14}^{\ast }$	92.47	43.66
$L_{VGG,IoU}$	${85.76}^{\ast }$	92.53	45.89
$L_{Dice,IoU}$	${85.13}^{\ast }$	92.61	44.90

Table 2. Table 2: B ​ F 1 𝐵 subscript 𝐹 1 BF_{1} values for different datasets and loss functions. For ISPRS dataset UNet model with ResNet34 backbone was used, for INRIA dataset – UNet with Inception-ResNet-v2 backbone

	AICD	ISPRS Potsdam	INRIA AIL
	ConvNN	ResNet34	Inc.-ResNet-v2
$L_{B{F}_1,IoU}$	71.51	82.48	51.75
$L_{IoU}$	68.01	81.02	50.75
$L_{SS,IoU}$	${20.74}^{\ast }$	81.55	29.53
$L_{VGG,IoU}$	${69.41}^{\ast }$	82.77	40.18
$L_{Dice,IoU}$	${70.49}^{\ast }$	81.86	29.45

Equations

$$w{L}_{B{F}_1}+(1-w)L_{IoU}$$

$$P^c=\frac{1}{|B_{pd}^c|}\sum _{x\in B_{pd}^c}[[d(x,B_{gt}^c)<\theta ]],R^c=\frac{1}{|B_{gt}^c|}\sum _{x\in B_{gt}^c}[[d(x,B_{pd}^c)<\theta ]],$$

$$B{F}_1^c=\frac{2P^c{R}^c}{P^c+R^c}.$$

$$y_{gt}^b=pool(1-y_{gt},{\theta }_0)-(1-y_{gt}),y_{pd}^b=pool(1-y_{pd},{\theta }_0)-(1-y_{pd}).$$

$$y_{gt}^{b,ext}=pool(y_{gt}^b,\theta ),y_{pd}^{b,ext}=pool(y_{pd}^b,\theta ).$$

$$P^c=\frac{sum(y_{pd}^b\circ y_{gt}^{b,ext})}{sum(y_{pd}^b)},R^c=\frac{sum(y_{gt}^b\circ y_{pd}^{b,ext})}{sum(y_{gt}^b)},$$

$$B{F}_1^c=\frac{2P^c{R}^c}{P^c+R^c},L_{B{F}_1^c}=1-B{F}_1^c.$$

Full text

1 Introduction

Semantic segmentation of remote sensing images is a critical process in the workflow of object-based image analysis, which aim is to assign each pixel to a semantic label [20, 13]. It has applications in environmental monitoring, urban planning, forestry, agriculture, and other geospatial analysis. Although a general image segmentation problem is relatively well investigated, we consider a particular type of this problem related to buildings segmentation.

The common aspect of urban aerial imagery is its high resolution (from 0.05 to 1.0 m). Higher GSD brings a lot of small details and structures, but also increases intra-class variance and decreases inter-class differences. This applies particularly to the class of buildings, including an enormous number of shapes and patterns. After the stage of segmenting buildings, further analysis could be done, such as distinct building classification and height estimation, population prediction, economic forecasting, etc. For most of the tasks, it is crucial to separate all buildings and minimize the number of false positive/false negative instances. Even with high-resolution imagery, it is difficult to segment every object separately, especially in high-density urban areas, although we get more accurate information about the boundaries of buildings. Therefore it is necessary to construct a method that increases the attention of a neural network, used for semantic segmentation, to borders of closely located objects.

CNNs for semantic segmentation typically use such loss functions as cross-entropy, dice score and direct $IoU$ optimization while training the network, but they are not sensitive enough to some misalignment of boundaries. Even being deviated from the ground truth by 5-10 pixels, predicted boundary does not significantly contribute to the value of the loss functions, mentioned above, and scores of common pixel-wise metrics. To better account for boundary pixels we select $BF_{1}$ metric (see original work [9] and its extension [12]) to construct a differentiable surrogate and use it in training. The surrogate is not used alone for training, but as a weighted sum with $IoU$ loss (from direct $IoU$ optimization). We found that the impact of the boundary component of the loss function should be gradually increased during training, and so we proposed a policy for the weight update. The assumption why this works is that for the first epochs network learns to label instances as shapeless blobs with very uncertain boundaries, and after finding an instance we only need to adjust its borders because they define the whole segments explicitly.

In experiments we also found that network trained with a combination of $IoU$ and boundary loss converges faster: we increase the confidence of the network in its predictions on the border as fast as for intra-segment pixels; with other losses, the mask near the borders is very blurred. The next point is that network better handles edge effects while baseline models tend to miss or distort masks on the edge of an input image. Besides, if it is highly important to separate neighboring buildings, methods of instance segmentation are often applied, which can require extra data preparation from raw masks and they are commonly multistage. However, binary segmentation with boundary loss is end-to-end. In experiments on the challenging datasets ISPRS Potsdam and INRIA AIL, we obtain $>93.8\%$ and $>74.3\%$ pixel-wise $IoU$ score respectively. There is also a comparison with another method for accurate delineation of curvilinear structures [17], which encounters difficulties in remote sensing tasks. Our results show a consistent increase in the performance for all models on both datasets in comparison to various loss functions.

The remainder of the paper is organized as follows. Section 2 briefly reviews related work. The proposed boundary loss, as well as the steps to construct its surrogate, are presented in Section 3. Experiments and results are discussed in Section 4, followed by concluding remarks in Section 5.

2 Related work

The subject of our work intersects with two branches of research, which are Deep Neural Networks (DNNs) and direct optimization of performance measures or their differentiable versions also called surrogates. Besides, we should mention methods for accurate boundary delineation, as one of them (see [17]) was compared to our approach.

DNNs and semantic segmentation. Semantic segmentation based on DNNs is a pixel-wise mapping to semantic labels. Since 2012 a lot of neural networks [26] for solving this task were constructed. The success of AlexNet [15] in ImageNet challenge marked the beginning of neural networks application in computer vision tasks such as classification, object detection, and semantic segmentation. Initially, for the latter problem, they adapted CNNs, used for classification. Namely, they remove fully connected layers, and use initial layers as an encoder to downsample an image and derive features that have stronger semantic information but low spatial resolution. Another part of the network called decoder obtains a high-resolution mask of labels, which upsamples features into the resulting mask.

For the last few years, DNNs have shown to be very effective for semantic segmentation task. There are several models now that are regarded as state-of-the-art. First of all, it is SegNet architecture [2] which fits online video segmentation very well. DeepLab network [6] uses atrous convolutions and CRFs at postprocessing step to refine small details of the segmentation. The FRRN [22] is an example of a model with multi-scale processing technique; it is based on two separate streams of data to better handle the high-level semantic information and the low-level pixel information simultaneously. Finally, UNet architecture [23] is very popular nowadays for its enormous flexibility. Among aforementioned models UNet remains preferable because of the high efficiency of feature extraction and ability to apply various neural network architectures as backbones. UNet is an asymmetric FCNN (Fully Convolutional Neural Network) with skip connections between the downsampling and upsampling paths. UNet has shown its high performance in competitions and research projects, and we use it for all experiments in the paper.

Performance measure optimization. There already exist papers concerning direct optimization of metrics or their differentiable versions. In various applications such measures as $IoU$, $F_{1}$-score, $ROC$-area, $mAP$ are widely used. Constructing surrogates of these measures is the most tricky part because it can be very hard or even impossible to replace non-differentiable operations with differentiable ones and keep computational efficiency at the same time. For the task of segmentation, $IoU$ is usually used to measure the performance of any segmentation approach. As a result, there exists a lot of its surrogates, and the goal is to minimize the gap between the actual $IoU$ value and its differentiable approximation. In paper [18] they proposed $NeuroIoU$ loss, which approximates naive $IoU$-loss with a neural network. Another approach [3] called Lovasz-Softmax loss, is based on the convex Lovasz extension of submodular losses. A critical property for this surrogate is that it effectively captures the absolute minimum of the original loss.

In this paper, we construct a differentiable version of the metric, presented in [9]. The motivation is that boundaries of segments explicitly define them and their extraction is necessary for accurate building segmentation. Authors of [9] proposed a novel boundary metric $BF_{1}$ that accounts for accuracy of contour extraction and overcomes the weaknesses of mainstream performance measures. For $BF_{1}$ we will construct a differentiable surrogate and show how it increases the accuracy of segmenting borders.

Below we provide a list of metrics and loss functions to be used for comparative analysis:

• •

Direct $IoU$ loss: $IoU=\frac{TP}{FP+TP+FN},$ $L_{IoU}=1-IoU$,

• •

$Dice$ loss: $F_{1}=\frac{2TP}{2TP+FN+FP},$ $L_{F_{1}}=1-F_{1}$,

• •

$SS$ loss: $SS=\lambda\frac{TN}{TN+FP}+(1-\lambda)\frac{TP}{TP+FN},$ $L_{SS}=1-SS,$ where $\lambda$ is a weight to balance two components,

• •

$VGG$ loss: implemented in [17].

The terms $TP,$ $FP,$ $TN,$ $FN$ denote pixel sets of true positive, false positive, true negative, false negative classes on a predicted binary map. For all loss functions, notations of corresponding metrics or approaches, inducing them, are provided in a subscript. By $L_{BF_{1},IoU}$ we denote a weighted sum

[TABLE]

for an arbitrary weight $w\in[0,1]$.

Accurate boundary delineation. One possible solution to extract borders accurately is to use CRF (Conditional Random Fields). It works as an extra layer above the output of the original neural network. CRF is applied to capture additional contextual information and to produce much more refined prediction [1]. Authors of [21] used CRF together with Bayesian decision theory and proposed a heuristic to maximize the value of $EIoEU$. Another approach, a combination of CRF and superpixels [29], takes advantage of superpixel edge information and the constraint relationship among different pixels. Proposed algorithm of boundary optimization made it possible to improve $IoU$ score by 3% on PASCAL VOC 2012 [11] and Cityscapes [8] datasets compared to the performance of plain FCN.

We compare an approach of [17] to the one, proposed in our paper. Authors claim that pixel-wise losses alone such as binary cross-entropy are unsuitable for predicting curvilinear structures. For this problem, they developed a new loss term based on features extracted with VGG19 network [27]. Conceptually their approach is similar to ours because we add our surrogate to the weighted conventional loss function. In contrast to VGG19 features, using our surrogate, we manually extract boundaries which are features too and encourage a neural network to draw attention to them much stronger.

3 Our Method

In this work, we use deep neural network UNet as a base model, because it is one of the most reliable and universal models. As we already pointed out in the introduction, there are a lot of different types of encoders and decoders that we can use as backbones. Of course, this network can be extended to two-headed version (UNet with two decoders), where the first head predicts whole segments and the other predicts only their boundaries. However, this approach requires additional computational resources and data preparation compared to the original UNet. Our interest is to propose a method that would not increase implementation and computational complexities. Further, we first introduce $BF_{1}$ metric from [9] and then describe how to construct its differentiable surrogate.

3.1 Boundary metric ($BF_{1}$)

Let $S_{gt}^{c}$, $S_{pd}^{c}$ be the binary maps of class $c$ in the ground truth and predicted segmentation respectively, $B_{gt}^{c}$, $B_{pd}^{c}$ — the boundaries for these binary maps. Then the precision and the recall for class $c$ are defined as follows:

[TABLE]

where brackets $[[\cdot]]$ denote indicator function of a logical expression, $d(\cdot)$ is Euclidean distance measured in pixels, $\theta$ - some predefined threshold on a distance; in all experiments we set $\theta$ to $3$ or $5$. Here distance is calculated from a point to one of two sets (ground truth or predicted boundary) as the shortest distance from the point to the point of the boundary. The $BF_{1}^{c}$ measure for class $c$ is defined as:

[TABLE]

3.2 Surrogate construction ($L_{BF_{1}}$)

To construct the differentiable version of the metric $BF_{1}$, we need to extract the boundary of any segment somehow. There are several ways to do this. Let us denote by $y_{pd}$ a binary map for an arbitrary class $c$ for some image $I$, predicted by a neural network, $y_{gt}$ a ground truth map for the same class and the same image. We expect values of the predicted map to be distributed in $[0,1]$, and values of the ground truth map to be in $\{0,1\}$. The boundaries can be defined as

[TABLE]

Here $pool(\cdot,\cdot)$ applies a pixel-wise max-pooling operation to the inverted predicted or ground truth binary map with a sliding window of size $\theta_{0}$ (hyperparameter $\theta_{0}$ must be as small as possible to extract vicious boundary; usually we set $\theta_{0}=3$). Value $(1-y_{gt,pd})$ corresponds to inversion of any pixel of the map. To compute Euclidean distances from pixels to boundaries a supporting map should be obtained, which is the map of the extended boundary:

[TABLE]

The value of hyperparameter $\theta$ can be determined as not greater than the minimum distance between neighboring segments of the binary ground truth map. After that precision and recall can be computed as follows:

[TABLE]

where operation $\circ$ denotes pixel-wise multiplication of two binary maps and operation $sum(\cdot)$ — pixel-wise summation of a binary map. Finally, the reconstructed metric and a corresponding loss function are defined as:

[TABLE]

All operations above are differentiable in terms of either derivative or subderivative. In Fig. 1 there is an example of how the proposed surrogate works. It is worth to mention that of course we can use convolutions with edge detection filters, such as Sobel, but in experiments, we found that this approach is not as effective as max-pooling. Nevertheless, both methods have their drawbacks. The disadvantages of both methods are that on the first step we extract boundary not narrower than $2$-$3$ pixels in width, but with Sobel filters, boundaries are even wider and more blurred. The main drawback of the used max-pooling operation is that gradients are passed only in tensor cells with a maximum value within a current sliding window. Sobel operators do not have this problem and take into account all pixels of the sliding window.

4 Experiments

4.1 Example on synthetic dataset (AICD)

To demonstrate capabilities of the novel boundary loss and test the hypothesis that $L_{BF_{1}}$ can assist $L_{IoU}$ a synthetic dataset was used. It consists of $800$ images with one primitive segment in each image. We trained on this dataset without augmentations and postprocessing a not deep fully convolutional neural network with four conv layers in the encoder and four conv layers in the decoder to show that $L_{BF_{1},IoU}$ outperforms other losses under simple conditions.

First of all, we present a working principle of the new loss in Fig. 1. In the upper row (see (b), (d) and (f)) we depict ground truth, where (b) is a binary mask for the ground truth segment, (d) is a binary mask of the gt boundary after applying (1). Then we compute expanded boundary (f) using pooling operations (2). The same pipeline is applied to the predicted mask (color intensity represents the probability that pixel belongs to a foreground). After that we obtain Precision map, which represents pixel-wise multiplication of maps corresponding to (d) and (g), and Recall map, which is a multiplication of maps (e) and (f). Precision and Recall are normalized pixel-wise sums of these maps.

Some examples of how $L_{BF_{1},IoU}$ loss outperforms $L_{IoU}$ loss are presented in Fig. 2. It is clear that in comparison with $L_{IoU}$ training with $L_{BF_{1},IoU}$ helps better delineate segment borders: they are very similar to ground truth boundaries and even meet at the corners in places where they should be. The second advantage is that the loss better handles edge effects. The third column clearly explains this kind of a problem when a building is located near the edge of the image. $L_{IoU}$ loss is unable to segment the building near the edge, however $L_{BF_{1},IoU}$ does it accurately. Finally, the boundary loss converges faster on this synthetic dataset, and such behavior is expected not only for this simple data distribution. The intuition behind is that network keeps attention mainly on a boundary, which consists of fewer pixels in comparison with the entire segment.

4.2 INRIA AIL dataset segmentation

The dataset [16] consists of 180 tiles (RGB) of 405 $km^{2}$ area in total with GSD of $0.3$ m. There are only two classes: building and not building.

Training on this dataset is more complex than on AICD. As for preprocessing, very strong augmentations were applied including vertical, horizontal flips, hue/saturation adjustments, many kinds of noises and blurs. Every orthophoto was split into patches of size $512\times 512$ and UNet with Inception-ResNet-v2 [30] as backbone was trained on these patches. After that, all predicted masks for every patch were merged into one mask of the whole orthophoto. Also, test-time augmentation (flips, reflections) was applied during the inference step. As for the training process, first, the network was trained from ImageNet [10] weights with a frozen encoder for three epochs, and the learning rate was set to $3\mathrm{e}{-3}$. After that all batch-normalization layers were unfrozen, and the learning rate was decreased to $1\mathrm{e}{-3}$. Then all layers were unfrozen and trained for 100 epochs with the learning rate of $1\mathrm{e}{-4}$. Finally, the network was trained for 60 epochs with the learning rate of $1\mathrm{e}{-5}$. We used Adam optimizer with default Keras [7] settings. There were several experiments with different loss functions. As for $L_{IoU}$, $L_{Dice}$, $L_{SS}$ and $L_{VGG}$ losses, they did not require any adjustments while training: these losses were trained as a weighted sum with $BCE$ one by one. As for $L_{BF_{1},IoU}$ loss, it requires an additional procedure for mini grid-search: after the 8th epoch for every 30 epochs and for every weight $w\in\{0.1,0.3,0.5,0.7,0.9\}$ in equation $(BCE+wL_{BF_{1}}+(1-w)L_{IoU})$ a network was trained. Then the best weight is chosen, and the process repeats. In Fig. 4 we see the evolution of $IoU$ metric while training with combinations of $BCE$ and $L_{IoU}$, $L_{Dice}$, $L_{SS}$, $L_{VGG}$, or $L_{BF_{1},IoU}$. $BF_{1}$ is also a very important metric for remote sensing image segmentation (see Fig. 4).

Finally, from figures, we see that the final segmentation is better than the baseline model trained with $L_{IoU}$ in terms of $IoU$ and $BF_{1}$. Now we want to understand what is better when training with the boundary loss. Below we provide an example comparing two segmentations.

In Fig. 6 on the right we see more accurate shapes of edges and corners. The important feature, which follows from Fig. 6, is that several buildings are standing in one row with a distance between adjacent constructions about 5 pixels. In comparison with $L_{IoU}$, $L_{BF_{1},IoU}$ our loss managed to separate instances of buildings much better.

4.3 ISPRS Potsdam dataset segmentation

The dataset [24] contains 38 patches, each consisting of a true orthophoto extracted from the larger mosaic. There are five channels (RGB+NIR+DSM) with GSD $0.2$ m for each channel. As a model, UNet was trained with very strong augmentations similar to that used for INRIA AIL. We used test-time augmentation (flips, reflections), set the learning rate to $1\mathrm{e}{-3}$ and divided it by a factor of $10$ if for $10$ epochs there was no improvement in validation loss. We trained a neural network on patches of orthophotos with Adam optimizer.

In Fig. 8 and Fig. 8 we see the evolution of $IoU$ and $BF_{1}$ metric while training with combinations of $BCE$ and $L_{IoU}$, $L_{Dice}$, $L_{SS}$, $L_{VGG}$, $L_{BF_{1},IoU}$. Here we see slightly faster convergence of $IoU$ for $L_{BF_{1},IoU}$, the final segmentation is 1-2% better than baseline $(BCE+L_{IoU})$. Also there is an advantage that training with $L_{BF_{1},IoU}$ loss is much less noisy.

Even though $IoU$ value is the best with $L_{BF_{1},IoU}$ training, boundary metric converges faster with VGG loss; authors of the original paper mentioned that this loss function is also efficient for boundaries delineation. In Fig. 9 and Fig. 10 there are several examples segmented with $L_{IoU}$ and the boundary loss.

Direct $IoU$ loss (Fig. 9) has poorer performance on edges of segments and complicated shapes, whereas boundary loss (Fig. 10) on the right retrieved almost all tricky shapes.

In Table 1 and Table 2 we provide final results of different losses in the task of binary segmentation of buildings for different models and datasets. Backbones for comparison in these tables are chosen is such way that they showed the best performance over all we have tried (VGG16, VGG19, ResNet34, DenseNet121, DenseNet169, Inception-ResNet-v2).

Finally, here are some constraints and recommendations:

It is better to use boundary loss only for segments with precise edges and corners such as buildings, roads etc. 2. 2.

Let us discuss the choice of parameters $\theta_{0}$ and $\theta$ (1) and (2) respectively. The best choice for the parameter $\theta_{0}$ is as less as possible, but at the same time it must be possible to extract solid boundaries of segments. The next constraint is that $\theta_{0}$ and $\theta$ should be smaller than minimum distance between segments, otherwise the performance of segmenting current element is affected by other segments because of overlapped expanded boundaries. Via the trial and error process we set $\theta_{0}$ to $3$ and $\theta$ to $5$-$7$ as a proper choice, because theses values deliver the most accurate boundaries in all experiments.

5 Conclusions

In this work, we introduced a novel loss function to encourage a neural network better taking into account segment boundaries. As motivation, we applied this approach to a synthetic dataset of binary segmentation and after that proved its validity on real-world data. Our experimental results demonstrate that optimizing $IoU$ assisted by $BF_{1}$ metric surrogate leads to better performance and more accurate boundaries delineation compared to using other loss functions. As for future work, we would like to extend our approach to handle multi-class semantic segmentation. We are going to elaborate on this approach by increasing its computational capabilities with large scale sparse convolutional neural networks [19], using a new loss function, specially tailored for imbalanced classification [28, 5], utilizing an approach for combining multi-modal data through CNNs features aggregation [4] and imposing a confidence measure on top of the segmentation model based on the non-parametric conformal approach [14, 25, 31].