Boundary Knowledge Translation based Reference Semantic Segmentation

1 Introduction

In recent years, deep neural networks have triumphed over many computer vision problems, including semantic segmentation, which is critical in emerging autonomous driving and medical image diagnostics applications. In general, training deep neural networks requires a humongous amount of labeled data, which is laborious and costly to collect and annotate. To alleviate the annotation burden, some learning techniques, such as few-shot learning and transfer learning, have been proposed. The former aims to train models using only a few annotated samples, while the latter focuses on transferring the models learned on one domain to another novel one. Despite the recent progress in few-shot and transfer learning, existing methods are still prone to either inferior results, or the rigorous requirement that the two tasks are strongly related and a large number of annotated samples are required.

∗Corresponding authors.

†These authors contributed equally to this work.

For many tasks, including few-shot and transfer learning, the performances of existing approaches, such as segmentation networks, deteriorate as the number of object categories increase, as demonstrated by prior works [Vinyals et al., 2016; Chen et al., 2019], and also by our experiments. The root cause lies in that, existing approaches are devised to recognize the category-wise features and segment the corresponding objects. Recently, boundary-aware features have been introduced to enhance the segmentation results, yet their core frameworks still focus on classifying category-wise features and segmenting the corresponding objects.

In the paper, we propose a novel Reference semantic segmentation Network (Ref-Net) based on visual boundary knowledge translation. In Ref-Net, only the target samples pass through the segmentation network, while the boundary knowledge of open-source datasets is translated into the segmentation network in an adversarial manner. That means only the data flow of the target dataset will pass through the segmentation network. Fig. 1 shows the differences between the proposed framework and related frameworks. More importantly, Ref-Net is the first proposed dataset-level knowledge translation framework, which is different from the modellevel knowledge transfer of related frameworks. The fatal limitation of model-level knowledge transfer is that the specific segmentation ability of the model for the source categories is beneficial for the segmentation of target categories, but also will limit the bound of the performance on the target categories. The most direct evidence is that the segmentation performance will drop dramatically when the number of the category increases. By contrast, it is noteworthy that the performance will increase when the number of the category increase for Ref-Net.

With a reference object, the human vision system can effortlessly find the same category objects in another image and precisely tell their visual boundary, even if they belong to an unknown category. Inspired by the above fact, the Ref-Net is devised with a Reference Segmentation Module (RSM) and a Boundary Knowledge Translation Module (BKTM), as shown in Fig. 2. In RSM, the reference object guides the segmentation network to find objects of the same category. Meanwhile, BKTM is proposed to assist the segmentation network in handling the accurate boundary segmentation, through translating the annotated boundary knowledge of open-source datasets. The object category of open-source datasets can be totally different from the target dataset.

RSM is devised for segmenting the same category objects in the target samples with several annotated reference objects. For the annotated reference image, the extracted features by the first branch are set as a condition, which will be concatenated with the target image features. Based on the condition, the semantic segmentation network learns to find and segment the same category objects in the target images.

To alleviate the burden of laborious and costly annotations, we adopt BKTM to translate the general boundary knowledge of abundant open-source segmentation datasets into the boundary segmentation ability of segmentation network. For accurate segmentation, the segmented object should not contain any background feature; meanwhile, the segmented background should not have residual object features. Inspired by this fact, BKTM is designed to comprise an inner boundary discriminator and an outer discriminator, as shown in Fig. 2. The outer boundary discriminator distinguishes whether the segmented objects contain the features of the outer background. Meanwhile, the inner boundary discriminator distinguishes whether the segmented background contains features of the inner objects.

Our contribution is therefore the first dataset-level knowledge translation based Ref-Net for semantic segmentation, which brings increased performance as the number of target category increases. Also, a boundary-aware self-supervision and a category-wise constraint are proposed to enhance the segmentation consistency on both the image-level and representation level, respectively. We evaluate the proposed RefNet on a wide domain of image datasets, and show that, with only ten annotated samples, our method achieves close results on par with fully supervised ones.

2 RelatedWorks

For boundary-aware semantic segmentation, the commonly adopted framework is a two-branch network that simultaneously predicts segmentation maps and boundaries [Takikawa et al., 2019]. Unlike predicting the boundary directly, some strategies, such as pixel's distance to boundary [Hayder et al., 2017], boundary-aware filtering [Khoreva et al., 2017], boundary refinement [Zhang et al., 2017], boundary weights [Qin et al., 2019], are proposed for improving the segmentation performance on boundary. Unlike the above methods, the proposed BKTM focuses on the outer boundary of the object and inner boundary of the background. Meanwhile, two boundary discriminators are devised for discriminating whether the outer boundary of the object and inner boundary of the background contain residual features.

GAN based segmentation contains two categories: mask distribution-based methods [Arbelle and Raviv, 2018] and composition fidelity based methods [Chen et al., 2019]. The former discriminates mask distribution between the predicted mask and GT mask, while the latter adopts discriminator to discriminate fidelity of the composite images. Unlike the above GAN-based methods, an adversarial strategy is adopted to translate the annotated boundary knowledge of the open-source dataset into the segmentation network by two boundary discriminators.

Few-shot Segmentation (FSS) aims at training a segmentation network that can segment new category well with only a few labeled samples of those new category. It contains parametric matching-based methods [Zhang et al., 2018; Xian et al., 2019], prototype-based methods [Wang et al., 2019a], GCN-based methods [Liu et al., 2021], R-CNN based methods [Yan et al., 2019] and proposal-free based methods [Gao et al., 2019]. Ref-Net is expected to gain a general segmentation ability and focus on the segmentation ability of the target categories. In addition, the category number of support category in FSS is usually larger than two, while Ref-Net can handle a single category source dataset.

Transfer learning based segmentation contains pseudosample generation methods [Han and Yin, 2017], iterative optimization methods [Zou et al., 2018], graph-based methods [Yang et al., 2020a], and distillation methods [Yang et al., 2020b]. Those methods aim at transferring the models' ability on source dataset into target datasets, where the capacity for source datasets still occupy the part of the segmentation ability of the model. However, the Ref-Net aims at learning a general segmentation ability with a reference image as a condition. Training on target datasets will bring more focused segmentation ability on the target categories. In addition, transfer learning based methods require the two domains as similar as possible, but Ref-Net has no such requirement.

3 ProposedMethod

The proposed Ref-Net is composed of a Reference Segmentation Module (RSM) and a Boundary Knowledge Translation Module (BKTM), as shown in Fig. 2. RSM is designed for segmenting the same category objects in the target samples with the reference image as condition, which is inspired by the human recognition mechanism with a reference object.

*

Reference segmentation module Reference Segmentation Module (RSM)

Boundary Knowledge Translation (BKT) module Boundary Knowledge Translation Module (BKTM)

Boundary-aware Self-supervision pixel-wise multiplication : : ∗∗ ∗ affine transformation matrix BKTM contains two boundary discriminators, which are devised for distinguishing whether segmented objects contain outer features of the background and segmented background contains inner features objects, respectively. BKTM aims to translate the annotated boundary knowledge of open-source datasets into the ability of the segmentation network, which will dramatically reduce the requirement of labeled samples for the target categories.

Formally, given a target dataset S with object category labels, a reference image dataset S with fine-grained object mask annotation and an open-source dataset O with finegrained object mask annotation, the goal of Ref-Net is to learn a segmentation network for S with S as condition and guidance, while BKTM translates the annotated boundary knowledge of O into the segmentation network. It is noticed that the reference image dataset S are chosen from the target dataset S, which has the same object category. The object category of open-source dataset O can be totally different from the category of the target dataset S.

3.1 Reference SegmentationModule

Inspired by the human recognition mechanism with a reference object, RSM is devised to be composed of two network branches: the reference feature extraction branch and the target segmentation branch. The target segmentation branch is designed to be an encoder-decoder architecture. To maintain the consistency of the feature space, the reference feature extraction branch F e θ has the same network architecture and parameter as the encoder of the target segmentation branch.

Give a target image x ∈ S and a reference image (xck ,mck ) ∈ S, the extracted representations can be denoted as r = F e θ (x) and rck = F e θ (xck ∗ mck ), where the ∗ denotes pixel-wise multiplication. With the concatenated representation [r, rck ], the decoder F d θ of the target segmentation branch predicts the mask ˜m = F d θ ([r, rck ]). For simplicity, the reference segmentation is formulated as ˜m = Fθ(x, xck ∗ mck ), where the Fθ denotes the two-branch segmentation network.

real triplet pseudo triplet fake triplet . Limited Sample Supervision. In the training stage, the limited annotated samples are also fed to the target segmentation branch, which will generate direct supervision information. So, given a random image (x,m) ∈ S and a reference image (xck ,mck ) ∈ S, the segmented result ˜m = Fθ(x,mck ∗ xck ) is expected to approximate the GT mask m, which can be achieved by minimizing the pixel-wise twoclass Dice loss Ldic:

where, τ is the Laplace smoothing parameter for preventing zero error and reducing overfitting, the GT mask m will be an all-zero mask when the random image x and reference image xck has no same category object. For the limited annotated samples, the data augmentation strategy is adopted to increase the samples' diversity.

Representation Consistency Constraint. The same category objects in the reference image and target image should have a similar representation distribution. So, we adopt the MaximumMean Discrepancy (MMD) to constrain the representation distribution consistency. With the reference image (xck ,mck ) ∈ S and the same category target image x, the predicted mask is ˜m = Fθ(x,mck ∗ xck ). The representation consistency loss Lrep is defined as follows:

where f is the feature encoder for reference image.

Boundary-aware Self-supervision. To reduce the number of labeled samples, inspired by [Wang et al., 2019b], we propose a boundary-aware self-supervision strategy, which can strengthen the boundary consistency of target objects. The core idea is that for the same network, the predicted mask of the transformed input image should be equal to the transformed mask predicted by the network with the original image as input. Formally, for the robust segmentation network, given an affine transformation matrix A, the segmented result ˜m′ = Fθ(Ax, xck ∗ mck ) of the transformed image Ax and the transformed result AFθ(x, xck ∗ mck ) should be consistent in the following way: Fθ(Ax, xck ∗ mck ) = AFθ(x, xck ∗ mck ). Furthermore, we obtain the boundary neighborhood weight map w′ and w as follows:

where, Dr and Er denote the dilation and erosion operation with a disk strel of radius r, respectively. The weight map w′ and w can strengthen the boundary consistency. The boundary-aware self-supervision loss Lsel is defined as follows:

where, w′ and w are the weight maps of the predict masks ˜m′ and ˜m, respectively. The boundary-aware self-supervision mechanism not only strengthens the boundary consistency but also eliminates the unreasonable holes.

3.2 Boundary Knowledge TranslationModule

Inspired by the fact that humans can segment the boundary of the object through distinguishing whether the inner and outer of boundary contain redundant features, we devise two boundary discriminators, which can translate the boundary knowledge of the open-source dataset O into the reference segmentation network.

Outer Boundary Discriminator. Randomly sampling a pair of samples xc ∈ S and (xck ,mck ) ∈ S from the target image set S and the reference image set S. Next, the segmentation network predicts the mask ˜m = Fθ(xc , xck ∗ mck ). Then, the segmented objects xco are computed by: xco = ˜m ∗ xc . The concatenated triplet Ioa = [xc , ˜m, xco ] is fed to the outer boundary discriminator Doφ , which discriminates whether the segmented objects xco contain the outer features of the background. In the paper, Ioa is regarded as a fake triplet. Meanwhile, choosing an annotated sample (x,m) ∈ O from the open-source dataset O, the corresponding Ioe = [x,m, xo ] is labeled as a real triplet. Furthermore, we reprocess the GT mask m of samples x by dilation operation Dr and get the pseudo triplet Ios = [x,Dr(m), x Do ], where x ∈ O and x Do = Dr(m) ∗ x. The generated pseudo triplet Ios will assist the outer boundary discriminator in distinguishing the outer features of the background. The adversarial optimization between the segmentation network and outer boundary discriminator will translate the outer boundary knowledge of the source dataset O into the segmentation network with the following outer boundary adversarial loss Lout adv :

where, the Poa , Pos , Poe are the segmented outer boundary distribution, pseudo outer boundary distribution, and real outer boundary distribution, respectively. The PIo is sampled uniformly along straight lines between pairs of points sampled from the distribution Poe and the segmentation network distribution Poa . The Io = εIoe +(1-ε)Ioa , where the ε is a random number between 0 and 1. The gradient penalty term is firstly proposed in WGAN-GP [Gulrajani et al., 2017]. The λ is the gradient penalty coefficient.

Inner Boundary Discriminator. The inner boundary discriminator Diϕ is devised for discriminating whether the segmented background contains the inner features of the object. To obtain the segmented background, the predict background mask ˜m′ and GT mask m′ are reprocessed with the Not-operation as follows: ˜m′ = [1] - ˜m, m′ = [1] - m, where the [1] denotes the unit matrix of m's size. Then, the corresponding fake triplet Iia = [x, ˜m′ , xb], real triplet Iie = [x,m′ , xb ] and pseudo triplet Iis = [x,Dr(m′ ), x Db ] are computed in the same manner as done in the outer boundary discriminator. The generated pseudo triplet Iis will also assist the inner boundary discriminator in distinguishing the inner features of objects. Similarly, the inner boundary adversarial loss Lin adv is defined as follows:

where, the Pia , Pis , Pie are the segmented inner boundary distribution, pseudo inner boundary distribution, and real inner boundary distribution. Ii = εIie +(1-ε)Iia . The optimization on Lin adv will translate the outer boundary knowledge of the open-source dataset into the segmentation network.

3.3 Complete Algorithm

To sum up, two boundary adversarial losses Lout adv and Lin adv are used to translate the visual boundary knowledge of the source dataset O into the segmentation network. The basic reconstruction loss Lrec is adopted to supervise the segmentation on tens of finely-grained annotated samples of reference image dataset S. The representation consistency loss Lrep and the self-supervision loss Lsel are devised for strengthening the category-wise representation consistency and the boundary-aware segmentation consistency on target datasets S. During training, we alternatively optimize the segmentation network Fθ and two boundary discriminators Doφ ,Diϕ using the randomly sampled samples from the reference image dataset S, target dataset S and the open-source dataset O, respectively. For training the segmentation network Fθ, the total loss Lseg is calculated as follows:

For training the inner and outer discriminators, the Eq.(5) and Eq.(6) are adopted.

4 Experiments

Dataset. The target datasets we adopted contain Cityscapes, SBD, THUR, Bird, Flower, Human. What's more, the open-source datasets (MSRA10K, MSRA-B, CSSD, ECSSD, DUT-OMRON, PASCAL-Context, HKU-IS, SOD, SIP1K) are merged into MixAll, which contains multiple categories. We do not use coarse data during training, due to our boundary loss which requires fine boundary annotation.

Network architecture. In the paper, the segmentation network we adopted is the DeeplabV3+ (backbone: resnet50 ) [Chen et al., 2017]. Some popular network architectures (Unet [Ronneberger et al., 2015], FPN [Lin et al., 2017], Linknet [Chaurasia and Culurciello, 2017], PSPNet [Zhao et al., 2017], PAN [Li et al., 2018]) are also tested.

Parameter setting. The parameters are set as follows: τ = 1, λ = 10, ξ = 1, ζ = 1, η = 1. In the generation of pseudo samples, the disk strel of radius r for the dilation and erosion operation is randomly sampled integer between 11 and 55. The interval iteration number between the segmentation network and discriminators is 5, the batch size is 64, Adam hyperparameters for two discriminators α = 0.0001, β1 = 0, β2 = 0.9. The learning rate for the segmentation network and two discriminators are set as 1e-4 .

Metric. The metrics we adopted include Pixel Accuracy (PA), Mean Pixel Accuracy (MPA), Mean Intersection over Union (MIoU), and Frequency Weighted Intersection over Union (FWIoU). Since the Dice index and IoU are positively correlated, the Dice index is omitted.

4.1 Comparing with SOTAMethods

In this section, the Ref-Net is compared with the SOTA methods, including unsupervised methods (CAC [Hsu et al., 2018], ReDO [Chen et al., 2019]), few-shot methods (SGOne [Zhang et al., 2018], PANet [Wang et al., 2019a], SPNet [Xian et al., 2019], CANet [Zhang et al., 2019]), weakly/semi-supervised methods (USSS [Kalluri et al., 2019], ALSSS [Hung et al., 2018]) and fully supervised methods (boundary-aware methods:{ Gated-SCNN [Takikawa et al., 2019], BFP [Ding et al., 2019]}, Unet [Ronneberger et al., 2015], FPN [Lin et al., 2017], LinkNet [Chaurasia and Culurciello, 2017], PSPNet [Zhao et al., 2017], PAN [Li et al.,

2018] and DeeplabV3+ [Chen et al., 2017]) on six datasets. For the semi-supervised methods, ten labeled samples are provided. Except for the panoramic, the target dataset and open-source dataset have no overlapped object category. For a fair comparison, the categories of each multiple category dataset (SBD and THUHR) are split into two non-overlapping parts. The fully supervised methods are trained with both two parts. The transfer learning based methods are initially trained on the half-category samples and then trained with specified labeled samples of the rest half-category. Table 1 shows the quantitative results, where we can see that most scores of R(10) achieve the state-of-the-art results on par with existing non-fully supervised methods. Even with more labeled samples, Trans.(100) only achieves higher scores than R(10) on THUR and Bird dataset. Moreover, with only 10 labeled samples, the Ref-Net can achieve better results than some fully supervised methods and close results on par with the best fully supervised method. Meanwhile, with fully supervised samples, the Ref-Net achieves higher scores on the complex dataset (Cityscapes and SBD), which demonstrates the advantage of Ref-Net for handling datasets with more categories. Note that the resolution of the Cityscapes dataset we adopted is 128 ∗ 128. The above two causes lead to the relatively low-scores of all the methods. However, it still validates the superior performance and wide application of the Ref-Net.

4.2 Results on Incremental Category Number

Table 2 gives the IoU scores of Ref-Net and DeepLabV3+ with the incremental category number. For the Ref-Net, the eight categories are set as target dataset, and ten labeled samples are provided for those categories. The rest of the categories are set as the open-source dataset. From Table 2, we can see that the IoU scores of bicycle increase with the incremental category number when the 'C-Num.' larger than 3. The first two scores (44.83 and 56.75) of the bicycle are larger than the score 30.85 ('C-Num.=3'). The reason is that the network trained with a single category will learn more category-aware features and focus on the single category. When the category number increases, the IoU scores first decrease then increase, which indicates that the Ref-Net can learn more general segmentation ability with the incremental category number. For the rest seven categories, most of the IoU scores increase with the incremental category number. In contrast, most IoU scores of DeeplabV3+ decrease with the incremental category number, which verifies the drawback of model-level knowledge translation.

4.3 Ablation Study

To verify each component's effectiveness, we conduct an ablation study on the boundary-aware self-supervision, the pseudo triplet, two boundary discriminators, supervised loss, and the different numbers of labeled samples. In the experiment, the knowledge translation is set as {THUR ← MixAll-}. For all R-, ten labeled samples of each category are provided. From Table 3, we can see that R(10) achieves higher scores than others, which verifies the effectiveness of each component. In addition, R(10) achieves about 2% increase on the scores of R-inner and R-outer and about 5% increase on the scores of R-disc , which demonstrates that the two-discriminator framework is useful for improving the segmentation performance. For the different numbers of labeled samples, we find that 10-labeled-samples is a critical cut-off point, which can supply relatively sufficient guidance. With more labeled samples, Ref-Net achieves better performance.

5 Conclusion

In this paper, we propose the Ref-Net for segmenting target objects with a reference image as condition, which is inspired by the human recognition mechanism with a reference object. The Ref-Net contains two modules: a Reference Segmentation Module (RSM) and a Boundary Knowledge Translation Module (BKTM). Given a reference object, RSM is trained for finding and segmenting the same category object with tens of finely-grained annotations. Meanwhile, MMD and boundary-aware self-supervision are introduced to constrain the representation consistency of the same category and boundary consistency of the segmentation mask, respectively. Furthermore, in BKTM, two boundary discriminators are devised for distinguishing whether the segmented objects and background contain residual features. BKTM is able to translate the boundary annotation of open-source samples into the segmentation network. The open-source dataset and target dataset can have totally different object categories. Exhaustive experiments demonstrate that the Ref-Net achieves close results on par with fully supervised methods on six datasets. The most important of all, the segmentation performance of all categories will improve with the increasing categories in the target dataset, while existing methods are opposite. The root reason is that the proposed method is based on datasetlevel knowledge translation, where the data flow of opensource samples will not pass the segmentation network. It brings a new perspective for the framework design of the image segmentation task.

Acknowledgments

This work is supported by National Natural Science Foundation of China (No.62002318), Zhejiang Provincial Natural Science Foundation of China (LQ21F020003), Key Research and Development Program of Zhejiang Province (2020C01023), Zhejiang Provincial Science and Technology Project for Public Welfare (LGF21F020020), Programs Supported by Ningbo Natural Science Foundation 202003N4318), Major Scientific Research Project of Zhejiang Lab (No.2019KD0AC01), and Zhejiang Lab (No.2020KD0AA06).