Edge-aware Plug-and-play Scheme for Semantic Segmentation

In semantic segmentation tasks, the higher the segmentation accuracy of the target, the more accurate the segmentation of its edge parts tends to be, and vice versa. Therefore, effectively incorporating edge supervision in the segmentation model can improve the segmentation accuracy of the network. Currently, there are two main approaches to incorporating edge supervision in segmentation models: parallel supervision and series supervision (as shown in Figure 1).

The parallel supervision usually adds an auxiliary head outside the backbone, which takes the original input image $x$ as its input and uses the edge label $\hat{y}_{e}$ for supervision to improve the segmentation accuracy of the decoder head on the backbone. For example, EG-CNN [6] uses edge gated layers to reconstruct the edges of the target, while ET-Net [34] uses the feature extraction ability of the first two layers of CNN to extract low-level features map for edge supervision. However, these methods are not suitable for ViT-based models that do not have multi-scale features. Chen et al. [1] combines edge detection tasks with segmentation tasks, and through a fusion network, fuses the results of edge detection and segmentation to improve the segmentation accuracy. Moreover, Li et al. [16] first extracts the edges from the segmentation results, and then uses edge and body supervision to separately segment the edges and bodies, and finally designs a fusion model to merge the results of both parts. However, the latter two methods require an additional fusion network to use edge supervision to improve segmentation accuracy. Therefore, the approach of extracting edges and segmentation results separately and then merging them is cumbersome.

The series supervision involves the attachment of an auxiliary edge detector to a segmentation network, with input provided by $\hat{y}$, the predicted output from the segmentation network. Through supervision using edge labels $\hat{y}_{e}$, the information extraction capabilities for edges in the backbone and decoder are influenced, resulting in improved segmentation accuracy. Zheng et al. proposed KLPPNet [38] that extracts edges through traditional algorithmic methods directly from the segmentation results and calculates its loss against edge labels $\hat{y}_{e}$. SEMEDA [4], on the other hand, introduced a segmentic edge detection network in series framework connected to the decoder head to exploit edge supervision information. However, such a series approach increases the architecture’s depth, which may trigger adverse effects on back-propagation outcomes during edge supervision and potentially lower the primary network’s edge information extraction capacity.

Most of the previous works above are designed with specific edge supervision modules for certain networks, whether in parallel or series supervision. However, such specific modules are difficult to transfer into other segmentation networks and their effectiveness cannot be guaranteed.

In semantic segmentation, the loss function plays a crucial role as it can significantly affect network learning efficiency. Existing segmentation losses have been classified into four categories by Ma et al. [20] and Jadon et al. [11]: distribution-based losses, region-based losses, compound losses, and boundary-based losses. Different types of loss functions have different optimization objectives and focuses, where distribution-based losses aim to improve overall classification accuracy, region-based losses aim to increase the overlap between predicted results and true labels, and compound losses combine the strengths of both types. The boundary-based loss, on the other hand, approaches it from a spatial geometry perspective by using the distance between the predicted and ground truth labels’ boundaries to construct the loss function. This study is centered on boundary-based loss functions, wherein some of the premier functions within this classification encompass the Boundary (BD) Loss [14] and the Hausdorff distance (HD) Loss [13].

BD Loss measures the quality of boundary prediction by calculating the distance between each pixel and the nearest boundary pixel in GT. The specific calculation formula can be defined as [14]:

where, $\Omega$ refers to the entire image area, $q\in\Omega$ is any pixel point on the image, $G\subseteq\Omega$ is the region where the GT exists, with binary pixel values of $\{0,1\}$, $S_{\theta}\subseteq\Omega$ is the predicted labeling area, with pixel values of $(0,1)$, and $D_{G}(q)$ is the distance between the pixel point $q$ and the nearest pixel points on the boundary of the region $G$.

HD Loss measures the accuracy of boundary prediction by computing the Hausdorff distance between predicted and GT boundaries. Its formula can be expressed as:

where, $A$ is the total number of pixels on the predicted boundary, $a_{i}$ represents a pixel on it, $B$ is the total number of pixels on the GT boundary, $b_{j}$ represents a pixel on it, and $d(a_{i},b_{j})$ is the Euclidean distance between pixel $a_{i}$ and $b_{j}$.

As both edge supervision and boundary-based losses are approached from a spatial-geometric perspective, their ideological origins are consistent. However, to the best of our knowledge, no research has emerged that have proposed an edge supervision scheme in combination with boundary-based losses, for semantic segmentation tasks.

To address the issue that the aforementioned relevant edge supervision methods are not easily applicable to other types of semantic segmentation networks, we propose an Edge-aware Plug-and-play Scheme (EPS). It is applicable to any semantic segmentation network and is easy to use. EPS mainly consists of two steps, which are to extract Edge GT with a thickness of $d_{e}$ and to copy decoder head.

Firstly, the edge thickness $d_{e}$ can reflect the degree of edge coarseness, and Edge GT with different edge thicknesses has different effects on edge supervision[38]. In order to generate Edge GT with a thickness of $d_{e}$, EPS uses the simplest edge extraction method by using a kernel of size $n\times n$ to process the GT image. The relationship between $d_{e}$ and $n$, as well as the calculation method of the kernel, are as follows:

Secondly, EPS copies the decoder head to generate an auxiliary head with a skip connection identical to the decoder head, but its weights are not shared with the decoder head. However, its supervision information is Edge GT (see Figure 2). It can be seen that EPS is an abstract strategy, and its specific implementation depends on the semantic segmentation network, so it has general applicability and is easy to implement, also plug-and-play. EPS only participates in the updating of model parameters during training and can be completely discarded during inference, without changing the size of the model during inference. Although this scheme is simple, our experiments have proven its effectiveness.

In EPS, we define the edge thickness $d_{e}$ of Edge GT as prior knowledge, and its value can be reflected by the size of kernel. For example, to set the edge thickness of Edge GT to $d_{e}=3$, a $7\times 7$ kernel is used to process GT. As the distribution-based loss optimizes globally, it doesn’t have the concept of edge thickness, so the edge segmentation thickness is random and uncertain. On the other hand, the extracted Edge GT in EPS has an explicit and equally-thick boundary with a thickness of $d_{e}$. To fully utilize this prior knowledge, we propose a boundary-based loss called Polar Hausdorff (PH) Loss. PH Loss calculates the Hausdorff distance between the internal and external edges in the predicted image $\hat{p}$ (as shown in Fig.3), and makes it tend to $d_{e}$, the edge thickness of Edge GT in EPS. This is different from HD Loss, which calculates the Hausdorff distance between the predicted image and GT.

Assuming the edge prediction image $\hat{p}$ is as shown in Fig. 3 (left), the edge image $\hat{p_{e}}$ with thickness 1 of the predicted edge image $\hat{p}$ is extracted (Fig. 3 right). According to EPS, the edge thickness of Edge GT is defined as a preset value $d{e}$. Therefore, during training, the Hausdorff distance between the inner edge pixel set $P$ and the outer edge pixel set $Q$ of the image in Figure 3 right should tend toward $d_{e}$, which is the optimization goal of PH Loss. Thus, PH Loss is specifically defined as:

where, $PHD_{P,Q}(\hat{p},n)$ refers to the Polar Hausdorff distance between internal and external edges, and the calculation of $d_{e}$ is shown in equation (6).

The computation of PH Loss is challenging in distinguishing the inner and outer edges of the edge image $\hat{p_{e}}$ for calculating the Polar Hausdorff distance between them since it is not feasible to classify them using one-pixel classifier, which can significantly increase the computational burden of the loss. To address this problem, we propose a more straightforward approach by utilizing the polar coordinate system to distinguish the inner and outer edges and compute their Hausdorff distance, as shown in Algorithm 1. Combining Fig.4, we first determine the geometrical center of the set $P\cup Q$ on $\hat{p_{e}}$, and then use centering to convert the coordinates of all pixels in $P$ and $Q$ into polar coordinates. We introduce a hyperparameter $n$ and draw $n$ rays from the geometric center outwards, with their angles uniformly starting from $0^{\circ}$ and increasing by $360^{\circ}/n$ each time. Then, for each angle $\theta$, we calculate the set of intersection points $P_{\theta}\cup Q_{\theta}$ of the ray with angle $\theta$ and the inner and outer edges, and compute the distance $\rho_{\theta}$ from each intersection point to the center. We select the minimum value $\rho_{min}$, and the corresponding intersection points $p$ lies on the inner edge intersection set $P_{\theta}$. We consider that the distances among the intersection points in $P_{\theta}$ are less than the threshold $\delta=2$, and select all the inner edge intersections $P_{\theta}$ via thresholding, while the remaining intersections $P_{\theta}\cup Q_{\theta}$ form the outer edge intersection set $Q_{\theta}$. Thus, we have distinguished the inner intersection sets $P_{\theta}$ and the outer intersections set $Q_{\theta}$ of the ray with angle $\theta$ between $P_{\theta}\cup Q_{\theta}$. By traversing each angle $\theta$ of the rays, finally, we calculate the Euclidean distance $d_{ph}(\theta)$ between the farthest points in $P_{\theta}$ and the nearest points in $Q_{\theta}$. The ultimate step is to select the maximum value from all of $\{d_{ph}(\theta)\}$ as the Polar Hausdorff distance $PHD_{P,Q}$.

The PH Loss is a proposed edge-supervision loss combined with EPS, which involves two hyperparameters. One is the hyperparameter for extracting the edge thickness $d_{e}$ of $\hat{p}_{e}$, which is also a hyperparameter in PH Loss. The other is the number of rays from the polar coordinate center in PH Loss, denoted as $n$. However, in subsequent experiments, it was found that $n$ had robusness, meaning that the choice of $n$ has little effect on the results. We recommend setting $n=8$.

Our proposed method was evaluated using the Cityscapes dataset, a well-known benchmark for semantic segmentation, which comprises 5000 high-resolution images with 19 categories. For hardware, we trained our models on a Linux server equipped with an I7 12700K CPU and two 24G RTX3090Ti GPUs. In terms of software framework, we utilized the PyTorch-based semantic segmentation framework MMSegmentation for all training and testing.

To ensure fair comparisons of experimental results, all models were trained with the official default parameters of MMSegmentation, such as the learning rate and momentum. For the loss function, we utilized the Cross-Entropy (CE) Loss as the base and combined it with the Polar Hausdorff (PH) Loss to achieve better segmentation accuracy. We reported segmentation accuracy using the standard mean Intersection over Union (mIoU) and mean Accuracy (mAcc) metrics.

To verify the versatility of EPS across different semantic segmentation models, we conducted experiments where our models were trained using Edge GT generated by a $5\times 5$ kernel. We assessed nearly all semantic segmentation models within MMSegmentation, including both those with and without auxiliary heads, such as CGNet, MobileNetV3, ANN, BiSeNetV2, and others. The specific models utilized and corresponding experimental outcomes are presented in Table 1. We observed a noticeable improvement in both mIoU and mAcc for almost all state-of-the-art models after implementing the EPS strategy. This signifies the compatibility of EPS with various semantic segmentation models and its plug-and-play nature, which enables its direct use without any consideration of the model’s unique attributes.

We selected four models, namely CGNet, ERFNet, SegFormer, and STDC, to further explore the experimental results of using PH Loss in EPS.The results are presented in Table 2. Based on our analysis, it can be inferred that the utilization of EPS has resulted in substantial enhancements in the performance of the SOTA. Furthermore, the incorporation of PH Loss has demonstrated a superior improvement in the model’s accuracy.

Our proposed PH Loss is designed to be used in conjunction with the EPS framework, with two hyperparameters: the edge thickness $d_{e}$ of the Edge GT and the number $n$ of candidate distances $d_{ph}$ in the $PHD_{P,Q}(\hat{p},n)$. In order to investigate the effects of these two hyperparameters, we conducted a series of ablation experiments on CGNet, ERFNet, SegFormer, and STDC.

After conducting an analysis of Table 3, we observe that there is no evident regularity for selecting the hyperparameter kernel, and the optimal kernel varies among different models. However, the performance is relatively favorable when selecting the kernel as $5\times 5$ and $7\times 7$. Additionally, employing EPS led to an improvement in mIoU, irrespective of the kernel size. Analysis of Table 4 revealed that even with PH Loss, the optimal kernel varied across different models. Comparing the results of the same kernel and model in Table 3, it was observed that almost all experiments using PH Loss performed better than those without PH Loss. In instances where the results of EPS with PH Loss were inferior to those of EPS with CE Loss, the cause may be attributed to insufficient training iterations. Finally, by analyzing Table 5, we found that the impact of selecting hyperparameter $n$ on the results is not significant, but as $n$ increases, the computational complexity of the model also increases. Consequently, it is recommended to select a smaller value for $n$, such as $n=8$.