TB-Net: A Three-Stream Boundary-Aware Network for Fine-Grained Pavement Disease Segmentation

The architecture of our TB-Net is shown in Figure 2. The three streams first capture different types of information, and the output features are then fused to obtain the segmentation map. The additional boundary map is produced in parallel to further refine the segmentation results.

Spatial Stream. Since encoding rich spatial information is critical in this dense prediction task, we follow [32] and stack three convolutional blocks for its effectiveness. Specifically, each block contains a convolutional layer using 3$\times$3 filters with a stride of 2, and the number of filters for each layer is [64,128,256]. A batch normalization and a ReLU activation are further applied. This shallow structure and the small stride have the benefits that allow the model to preserve the spatial size of the original image, as well as retaining affluent spatial information.

Context Stream. Capturing contextual dependencies plays an important role in semantic segmentation [34, 12, 32], thus we propose a context stream that can model a wide range of contextual relationships over local features in order to achieve better feature representations. Given the input image, it incorporates the multi-scale contexts by utilizing a context-aware attention (CAA) mechanism (detailed in Section 2.2). We first make use of the ImageNet-pretrained ResNet-101 v2 [17] to obtain the features from the last layers of the second and the last residual blocks, respectively. In this way, the middle-level and the high-level feature extraction can be achieved. After applying a 1$\times$1 convolutional layer, the transformed features are then fed into two CAA modules to encode the global contextual information and generate two enhanced features.

Boundary Stream. Making use of boundary information can enhance the segmentation performance, particularly for thin and small objects [8, 27]. In order to refine the fine-grained segmentation results, we propose a boundary stream that takes the learned global contextual features as the input and generates boundary predictions. This allows the information to flow from the context to the boundary stream. Specifically, we first utilize a residual block with a short-skip connection [14, 38]. This aims to force the network to pay more attention to the informative features. Inspired by the recent success of gated convolutions that learn dynamic feature selection for each channel and each spatial location on images [33, 19], we then apply a global-gated convolution module (detailed in Section 2.3) that assigns more weights to boundary-related pixels by incorporating the global contexts. The output boundary predictions are obtained by further utilizing two transposed convolutional blocks followed by a sigmoid activation.

Feature Fusion. In the feature fusion module, different types of features are fused to produce the refined segmentation outputs. We first concatenate and transform the outputs of the spatial and the context streams by a batch normalization and a convolutional layer. The encoded features are then fused with the boundary information. Similar to [18], the learned feature maps are refined by using a reweighting mechanism to generate a weight vector. After that, a matrix multiplication operation is performed between the weight vector and the input features, and an element-wise sum operation is further applied to obtain the learned features. The final features are resized to the original image using bilinear interpolation to produce the segmentation map.

Attention mechanism [29, 12, 22] allows the input to interact with each other and models the global dependencies. Similar to [12] where the self-attention is utilized to capture the semantic interdependencies between any two positions, we employ context-aware attention modules in the context stream to enrich feature representations.

Let the input feature map be $F_{c}$. As illustrated in Figure 3(a), three convolutional layers are used to generate the new features $F_{c}^{1}$, $F_{c}^{2}$ and $F_{c}^{3}$. Then the matrix multiplication is utilized on $F_{c}^{1}$ and $F_{c}^{2}$ after a reshaping operation followed by a sigmoid activation to obtain an attention map. Next, another matrix multiplication is performed between $F_{c}^{3}$ and the attention map. The learned context-aware feature $\hat{F}_{c}$ can be obtained by using an element-wise sum operation with a scale parameter $\gamma$, which is formulated as:

$$\hat{F}_{c}=\gamma(F_{c}^{3}\otimes\sigma(F_{c}^{1}\otimes(F_{c}^{2})^{T}))\oplus F_{c},$$

(1)

where $\sigma$ is the nonlinear activation function, and $\otimes$ and $\oplus$ denote the matrix multiplication and the element-wise sum operation, respectively.

Here the learnable $\gamma$ is initialized as 0, which allows the network to rely on the local cues before assigning more weights to the non-local evidence [35]. In this way, the context-aware attention modules enable the model to capture the affluent relationships by incorporating the semantic interdependencies of the inputs. Note that since there are two context-aware attention modules used in the context stream as shown in Figure 2, the learnable scale parameters in these two modules are defined as $\gamma^{1}$ and $\gamma^{2}$ separately.

Inspired by [27] where a gated mechanism is proposed to ensure that only the boundary-relevant information is processed, we introduce the global-gated convolution module. This aims to highlight the relevant information and filter out the rest to better generate the boundary representations and further improve the segmentation performance. The architecture of the GGC module is shown in Figure 3(b).

Specifically, let the two input features be $\tilde{F}_{c}^{1}$ and $\tilde{F}_{c}^{2}$. We first concatenate these two features, then perform a batch normalization to balance the different scales of the input features. Next, two 3$\times$3 convolutional layers using 512 filters are applied with a ReLU in between, followed by another batch normalization and a sigmoid activation. After that we multiply the results with the input feature $\tilde{F}_{c}^{1}$ and perform an element-wise sum operation to produce the boundary features $F_{b}$. In this way, the model can learn to give the boundary regions more weights. The overall process can be summarized as:

$$F_{b}=\tilde{F}_{c}^{1}\otimes\sigma(\delta(\tilde{F}_{c}^{1}||\tilde{F}_{c}^{2}))\oplus\tilde{F}_{c}^{1},$$

(2)

where $\delta$ denotes the normalized 1$\times$1 convolutional layers and $||$ is the concatenation operation. $\sigma$, $\otimes$ and $\oplus$ are the same as in Equation (1).

The proposed TB-Net is trained using a dual-task loss on the predicted segmentation and boundary maps. For the segmentation prediction, a cross-entropy loss with a softmax activation is applied on the predicted outputs of the feature fusion module. Due to the fact that pixels in the minority classes tend to be less well trained, we further utilize a weighting mechanism for the highly imbalanced fine-grained disease data, which can be written as:

$$\mathcal{L}_{WCE}=-\sum_{n=1}^{N}\omega_{n}\sum_{c=1}^{C}\sum_{m=1}^{|\mathbf{I}_{n}|}y_{m,c}^{(n)}\cdot\log(p_{m,c}^{(n)}),$$

(3)

where $\mathcal{L}_{WCE}$ is the loss on the predicted segmentation map, $N$ denotes the total number of the training images, and $|\mathbf{I}_{n}|$ is the number of pixels in each sample. $y_{m,c}^{(n)}$ is the label indicating whether the given pixel in position $m$ of the image $n$ belongs to class $c$ out of the total $C$ classes, and $p_{m,c}^{(n)}$ is the softmax probability of the corresponding predicted output of the network. The weight $\omega_{n}$ for each image can be defined as:

$$\omega_{n}=\sum_{c=1}^{C}\sum_{m=1}^{|\mathbf{I}_{n}|}\omega_{c}\cdot y_{m,c}^{(n)},$$

(4)

where $\omega_{c}$ is the class weight which can be computed over all the training images based on the number of pixels for each class by:

$$\tilde{\omega}_{c}=\sum_{c=1}^{C}\sum_{n=1}^{N}\sum_{m=1}^{|\mathbf{I}_{n}|}y_{m,c}^{(n)}\Big{/}\sum_{n=1}^{N}\sum_{m=1}^{|\mathbf{I}_{n}|}y_{m,c}^{(n)}.$$

(5)

We then normalize the weight by the sum of $\tilde{\omega}_{c}$ and obtain the final class weight $\omega_{c}$. In this way, the weighting mechanism penalizes the loss based on the estimated labels. This is beneficial to reduce the class imbalance that is inherent in the data and avoid ignoring the minority classes.

For the boundary prediction, a binary cross-entropy loss is utilized by introducing the ground-truth boundary information. Here we make use of a recent state-of-the-art method on edge detection [15] and generate binary ground-truth boundaries given the annotated segmentation maps. For each image, let $\bar{y}_{m}^{(n)}\in\{0,1\}$ be the pixel-wise label indicating whether it is a boundary pixel, and $\bar{p}_{m}^{(n)}$ is the corresponding prediction. The binary cross-entropy loss $\mathcal{L}_{BCE}$ on the predicted maps can be written as:

$$\mathcal{L}_{BCE}=-\sum_{n=1}^{N}\sum_{m=1}^{|\mathbf{I}_{n}|}(\bar{y}_{m}^{(n)}\cdot\log(\bar{p}_{m}^{(n)})+(1-\bar{y}_{m}^{(n)})\cdot\log(1-\bar{p}_{m}^{(n)})).$$

(6)

The segmentation and the boundary predictions are jointly learned using the proposed dual-task loss $\mathcal{L}$ as follows:

$$\mathcal{L}=\lambda_{1}\mathcal{L}_{WCE}+\lambda_{2}\mathcal{L}_{BCE},$$

(7)

where $\lambda_{1}$ and $\lambda_{2}$ are the balancing parameters.

To evaluate the proposed method, we conduct the experiments on an airport pavement disease dataset which specifically targets this fine-grained task to reflect the road surface condition. The dataset contains 3946 images captured with a gray-scale camera on the complex cement/asphalt pavements under different illumination. The examples can be seen in Figure 1. In total, there are 3171 images for training, and 793 images for testing, which involve eight semantic classes and one background class with pixel-level labels. We randomly select 20 images from the training set for validation. The size of each image ranges from 640$\times$415 to 1800$\times$900, and the number of the annotated areas in the dataset for each class is: $\sharp$Crack: 3586, $\sharp$Cornerfracture: 151, $\sharp$Seambroken: 557, $\sharp$Patch: 312, $\sharp$Repair: 893, $\sharp$Slab: 3040, $\sharp$Track: 3749, $\sharp$Light: 58.

We evaluate our performance using two commonly used metrics [24] in semantic segmentation: Class Pixel Accuracy (CPA) and Intersection-over-Union (IoU) which are defined as:

$\displaystyle\begin{split}&CPA=\frac{p_{ii}}{\sum_{j=1}^{C}p_{ij}},\\ &IoU=\frac{p_{ii}}{\sum_{j=1}^{C}p_{ij}+\sum_{i=1}^{C}p_{ji}-p_{ii}},\end{split}$

(8)

where $p_{ij}$ is the number of pixels of class $i$ predicted to belong to class $j$ out of the total $C$ classes.

Our TB-Net is implemented using Tensorflow [1] on a single NVIDIA Tesla V100 16GB GPU. We adopt RMSProp optimizer [28] for optimizing the models. In the training stage, the initial learning rate is set to 0.0001, the decay is set to 0.995 and the model gets converged after 150 epochs. The input images and the corresponding ground-truths are resized uniformly to 512$\times$512. In the two context-aware attention modules, since the pre-activation ResNet does not have batch normalization or ReLU in the residual unit output, we further apply a batch normalization and a ReLU on the extracted feature maps. The learned boundary features are upsampled by bilinear interpolation and transformed by a 1$\times$1 convolution before being concatenated in the feature fusion module. We experimentally set the balancing parameters $\lambda_{1}=\lambda_{2}=1$. The overall training time is 35 hours and inference is 0.03 seconds.

We compare the performance of our TB-Net against four segmentation methods that attain the state-of-the-art results in other semantic segmentation tasks and can be potentially used for this fine-grained pavement segmentation task: (1) DeepLabv3 [6], a network employs atrous convolution with upsampled filters to capture the multi-scale contexts; (2) DeepLabv3+ [7], an encoder-decoder-based model that extends DeepLabv3 by adding a decoder to further refine the segmentation results; (3) DenseASPP [31], a method concatenates different atrous-convolved features with the multiple scales; (4) BiSeNet [32], a network combines two paths to achieve the rich spatial information and the sizeable receptive field. Note that we keep the learning objective functions the same for fair comparisons.

The results are shown in Table 1 and 2, where BiSeNet+ and BiSeNet are the same, except that BiSeNet+ does not make use of the weighting mechanism in the loss function. It can be observed that the proposed TB-Net consistently outperforms the existing methods in terms of both mean CPA (mPA) and mean IoU (mIoU). Specifically, for Cornerfracture and Seambroken where the numbers of the annotated areas are small in the dataset ($\sharp$Cornerfracture: 151, $\sharp$Seambroken: 557) with each annotated area being relatively small, we still achieve favorable performance over other competing methods. This indicates the importance of fusing the rich spatial and the contextual features with additional boundary information, which can effectively enhance the feature discrimination.

It is worth noting that Track contains the markings of various curves and straight lines as well as the low contrast water/oil stains, and our model improves the performance by large margins (56.78% vs. 43.29% in terms of CPA and 42.19% vs. 31.28% in terms of IoU). The overall results show that the proposed TB-Net can bring great benefits to fine-grained segmentation across different diseases with varying patterns.

We provide visual comparisons of the results obtained by our proposed TB-Net and one of the competing models, BiSeNet [32]. We also show the boundaries obtained from the proposed TB-Net and the corresponding ground-truths that are obtained using [15]. As shown in Figure 4, we observe that our model generally achieves better pixel-wise disease segmentation. Specifically, for Crack and Seambroken where the size of each disease tends to be small, more pixels that belong to these two classes are properly detected by our method. For other diseases where the sizes tend to be large, such as Patch and Repair, our model can also produce more robust and complete predictions.

Moreover, we can see that the predicted boundaries decently outline the different pavement diseases and help the network produce sharp segmentation results. For example, in the first row in Figure 4, the region for Seambroken is rather small and thin which poses more challenges, and BiSeNet fails to predict it correctly. However, our model that incorporates the boundary features better segments this type of disease and produces a higher-quality boundary.

The effect of the context-aware attention module. To investigate the use of the context-aware attention in our proposed TB-Net, we drop the attention module in the context stream and compare the performance to the full model. The results are shown in Table 3. Here we use TB-Net w/o attn to denote the model that does not make use of the context-aware attention module. From the table, we observe that the attention module generally helps improve the performance, especially in Track with different patterns and sizes of markings and stains, bringing 20.15% improvements in terms of CPA and 14.55% in terms of IoU. This indicates that the context-aware attention module can effectively allow the model to incorporate the long-range dependencies between different pixels and further enhance feature representations by capturing the contextual information.

The effect of incorporating the boundary information. We further perform an ablation study to analyze the effect of the boundary information. Table 4 illustrates the performance when dropping the boundary stream (denoted as TB-Net w/o BS). Here we fuse the encoded spatial and contextual features and apply a convolutional block, that consists of a 3$\times$3 convolutional layer, a batch normalization and a ReLU. The learned features are then taken as the input to the reweighting mechanism and further predict the segmentation results. We observe that the overall CPA and IoU drop or remain relatively unchanged after dropping the boundary stream for different classes except Patch. We hypothesize that this is due to the fact that it is less dependent on the boundaries to segment rectangle-like cement patches. For other small diseases such as Crack and Seambroken which contain much detailed information, it tends to be more beneficial from fusing the boundary stream.

The effect of the weighting mechanism in the loss function. Since the fine-grained disease data we use is highly imbalanced, we adopt a weighting mechanism in the loss function to avoid the model ignoring the minority classes. The effect of introducing the weighting mechanism can be seen in Table 5. Here we use TB-Net w/o w to denote the model that does not make use of the weighting mechanism. We observe that the full model significantly outperforms the model without the weighting mechanism in terms of both mPA and mIoU. Especially in Crack where the number of pixels in each annotated area tends to be small, the performance improves by large margins (45.24% and 10.48% in terms of CPA and IoU) by using the weighting mechanism. Besides, the performance of BiSeNet increases compared to BiSeNet+ in Table 1, which also demonstrates the effectiveness of this weighting mechanism.