Semi-supervised Semantic Segmentation using Auxiliary Network

1 INTRODUCTION . Semantic Segmentation , which identifies the category label of each pixel , is an important task in computer vision . A number of convolutional neural network ( CNN ) based semantic segmentation systems have been developed in recent years such as Chen et al . ( 2017a ; b ; 2018 ) ; Long et al . ( 2015 ) ; Zhao et al . ( 2017 ) ; Ronneberger et al . ( 2015 ) . For practical applications such as autonomous driving , there is a high demand for real-time processing and sufficient training data . Hence , different research directions have been explored . For example , some researchers proposed various efficient network structures ( Chen et al. , 2019 ; Paszke et al. , 2016 ; Poudel et al. , 2018 ; Lo et al. , 2018 ; Poudel et al. , 2019 ; Yu et al. , 2018 ; Zhao et al. , 2018 ; Romera et al. , 2017 ; Li et al. , 2019 ) , and the others focus on the weakly- or semi-supervised learning schemes ( Hung et al. , 2018 ; Lin et al. , 2016 ; Bearman et al. , 2016 ; Qi et al. , 2016 ; Rajchl et al. , 2016 ; Pathak et al. , 2015 ; Papandreou et al. , 2015 ; Hong et al. , 2015 ; Luc et al. , 2016 ; Dai et al. , 2015 ; Pathak et al. , 2014 ; Zhu et al. , 2019 ; Mittal et al. , 2019 ) . In this study , we first design an efficient segmentation network inherited from Chen et al . ( 2019 ) as our baseline model , and then propose a semi-supervised learning scheme . There are three training stages in our system . In the first stage , the segmentation network is trained using the labeled images . Then , in the second stage , an auxiliary network is trained using the segmentation network model trained in the previous stage and the labeled data to generate the confidence map . Inspired by Hung et al . ( 2018 ) , we use the concept of confidence map to assign a confidence score to each segmented pixel on the unlabeled images . In their work ( Hung et al. , 2018 ) , the confidence map is the output of a GAN discriminator network ( Goodfellow et al. , 2014 ) , where the discriminator network learns to distinguish between the segmentation map and the ground truth map . However , we find that the trusted ( high confidence ) regions in their confidence map are mostly located on the large target objects , and thus the effectiveness of semi-supervised learning is limited . To obtain more reliable small object labels , we adopt a different approach in generating the confidence map . In our approach , the confidence map is generated by an auxiliary ( CNN ) network , which is trained using the proposed auxiliary loss function . We carefully design the auxiliary loss function such that it leads to a reliable confidence map , particularly , on the small objects . In the third stage , the unlabeled images are used as inputs to the proposed system to generate the labels . Therefore , we can use both originally labeled and unlabeled images to retrain the segmentation network model to achieve a better performance in the end . In summary , our main contributions of this work are as follows . First , Based on DSNet-fast ( Chen et al. , 2019 ) , we design a powerful segmentation network , which achieves a very good balance between speed and accuracy . It produces 73.9 % mean IoU on the Cityscapes testing set with a speed of 73 FPS on a single GTX 1080Ti card . Second , We propose a semi-supervised learning scheme with the notion of auxiliary network , which can be used to annotate the unlabeled images . Third , with the help of auxiliary network , our semi-supervised learning method can include the unlabeled images in training , which improves the segmentation accuracy . 2 RELATED WORK . Recently , CNNs have been widely used in many fields of computer vision . For the semantic segmentation task , FCN ( Long et al. , 2015 ) is a pioneer . It replaces the fully-connected layers of the classification network ( Krizhevsky et al. , 2012 ) with the convolution layers , and thus it can generate dense labels on the input image of the same size . Then , SegNet ( Badrinarayanan et al. , 2017 ) was subsequently proposed , which uses a symmetric encoder-decoder architecture for feature map down-sampling and up-sampling . The U-Net ( Ronneberger et al. , 2015 ) introduces the concatenation operation to up-sample the features with different levels . PSPNet ( Zhao et al. , 2017 ) and DeepLab ( Chen et al. , 2017a ; b ; 2018 ) propose the atrous spatial pyramid pooling ( ASPP ) module to integrate multi-scale features . There are many other studies on improving the segmentation results ; however , for practical applications , a high inference speed network is very desirable . For example , ENet ( Paszke et al. , 2016 ) is a real-time segmentation model with good segmentation results . ICNet ( Zhao et al. , 2018 ) and BiSeNet ( Yu et al. , 2018 ) were recently proposed , and they aim at a better balance between speed and accuracy . Another challenge in training a semantic segmentation network is the need of a large amount of labeled data ( ground truth ) . The labeling cost of pixel-level annotation is extremely expensive . Hence , in recent years , the study of the weakly- and semi-supervised learning approaches to tackle this problem attracted a lot of attention . For the weakly supervised methods , the pixel-level annotation is no longer used to train the segmentation network , but instead , other forms of annotations obtained at low costs are used for training . Several different types of weakly-supervision have been studied , including bounding box supervision ( Rajchl et al. , 2016 ; Dai et al. , 2015 ) , scribbles supervision ( Lin et al. , 2016 ) , point supervision ( Bearman et al. , 2016 ) , and image-level supervision ( Qi et al. , 2016 ; Pathak et al. , 2014 ) . Particularly , the development of image-level supervision is the most popular one . Pathak et al . ( 2015 ) convert image-level labels to restrict the distribution of CNN output . Papandreou et al . ( 2015 ) combine image-level labels using EM algorithms to train the segmentation model . Wei et al . ( 2018 ) propose a generic classification network , which adopts the convolutional blocks with different dilated rates to generate dense object localization maps . It can produce reliable segmentation masks for training the segmentation model . For semi-supervised learning , Hong et al . ( 2015 ) propose a method to separately train the classification and segmentation networks , and then pass the information from the classification network to the segmentation network to reduce the search space for effective segmentation . Luc et al . ( 2016 ) employ an adversarial network to enhance the segmentation quality . Recently , Zhu et al . ( 2019 ) propose a video prediction model to generate new samples for improving segmentation accuracy . Moreover Hung et al . ( 2018 ) propose a self-taught learning approach based on the dense pixel-level probability maps produced by adversarial network . Mittal et al . ( 2019 ) describe the dual-branch method including a semi-supervised classification branch to learn from unlabeled data . In this study , we propose an auxiliary network , which can determine the credibility of each pixel on the segmentation map of unlabeled images and then the trusted pixels can be used as additional ground truth data in training . Different from the previous works , we focus on the quality of the generated labels . We believe that the supervision signals of small objects are more important than the larger ones . Thus , we assign different weights to the generated tags of different categories , especially biased towards the small objects . In training the auxiliary network , we propose an auxiliary loss function containing two critical terms to achieve this objective . 3 METHODOLOGY . In this section , we first describe the general framework of our semi-supervised semantic segmentation scheme , and then we describe the architecture and design concepts of the proposed segmentation network and auxiliary network . 3.1 SEMI-SUPERVISED LEARNING USING AUXILIARY NETWORK . There are in total three training stages in our semi-supervised learning scheme . In the first stage , we simply train the segmentation network using the labeled images . In this paper , we simply use the typical cross entropy loss in the supervised training of segmentation network . Then in the second stage , an auxiliary network is trained using the proposed auxiliary loss function . Both the originally labeled images and the previously trained segmentation network are used in this stage . Our auxiliary network takes the class probability maps of size H×W ×C produced by the segmentation network as input , where H and W are the height and width of the image size and C denotes the number of classes . The auxiliary network outputs a confidence map of size H ×W × 1 . Each pixel of the confidence map is a probability value between 0 and 1 , which represents the credibility of that pixel on the segmentation map . In the third stage , we feed the unlabeled images into the segmentation network pre-trained in the first stage to obtain the class probability map , and then pass this map through the auxiliary network to obtain the confidence map . Based on the probability value of each pixel of the confidence map , we can determine the segmentation reliability of each pixel on the segmentation map . Then , we select the pixels with high reliability as the annotated data points . And we ignore the pixels with lower reliability in the next-phase training . We set a threshold Taux on the confidence map to separate the high reliability pixels from the low reliability pixels . Hence , we can mask the error-prone pixels on the output segmentation map and produce the reliable annotated maps for the unlabeled images . Finally , since all the images now have labels , we retrain the segmentation network using all of them , and produce a more accurate segmentation model . Figure 1 shows the process for generating annotated labels for the images without ground truth . Loss function for auxiliary network To train the auxiliary network , we create another auxiliary ground truth map of dimension H ×W × 1 . Assuming that we have the original labeled ground truth class map and the estimated ( segmented ) class map produced by the segmentation network , the pixel value on the auxiliary ground truth map is set to 1 if the class at that pixel on the estimated segmentation map matches that on the ground truth class map ; and it is set to 0 , if not . Different from the typical binary cross entropy loss , we assign weights to the terms in the loss function for auxiliary network , which is defined below . Laux = − ∑ h , w Wnh , w ( γ n h , w × ynh , w log ( ah , w ( xn ) ) + ( 1− ynh , w ) log ( 1− ah , w ( xn ) ) ) ( 1 ) where ynh , w equals to 1 if the estimate class on the segmentation map at pixel ( h , w ) is identical to the class value on the ground truth label map at the same pixel , and ynh , w equals to 0 if they are different . In equation ( 1 ) , ah , w ( xn ) denotes the confidence map output at pixel ( h , w ) , when the auxiliary network takes the class probability map xn as input . Wnh , w represents the class weighting value at pixel ( h , w ) on the confidence map , and it is defined as 1/ ln ( c + pclass ) , inspired by ENet ( Paszke et al. , 2016 ) , where we set c equal to 1.04 . The class weighting at different locations of confidence map is decided by the ground truth label map . For example , when the class at the pixel ( h , w ) on ground truth label map is class i , the class weight at that pixel on the confidence map is the weight value associated with class i . Class weighting term is used to tackle the data imbalance issue . That is , some classes having few number of pixels are often unfavored in training . With class weighting , our auxiliary network performs well not only on the big objects but also on the small objects . Moreover , for different classes , the imbalance issue between positive and negative samples are quite different . Thus , we propose a weight term γnh , w , which is defined as # negative pixelclass/ # positive pixelclass , and it is calculated using the auxiliary network ground truth . Similarly , this weight term is location variant depending on the ground truth label map . Both Wnh , w and γ n h , w are the critical terms in the proposed auxiliary loss function . There were several versions tested and the current version gives the best result .

Summary: - This paper proposes a semi-supervised learning strategy for semantic segmentation of road scenes. Specifically, authors propose to include an auxiliary network that will predict the confidence (at pixel-level) of the predictions on unlabeled images. These confidence values will be used to  generate a new auxiliary ground-truth to retrain the network using the unlabeled images.