Coupled Segmentation and Edge Learning via Dynamic Graph Propagation

1 Introduction . Image segmentation and edge detection have been widely studied as important perception problems . The two problems are closely related . In fact , segmentation subsumes edge detection since any segmentation contour makes a closed boundary of a region . The converse is however not true since general edges do not always form closed contours . Nevertheless , edge detection can serve as an auxiliary task to improve segmentation performance since edges provide important pairwise similarity cues for segmentation . Early works tend to focus on the grouping and contrast of pixels from a perceptual similarity perspective . Martin et al . [ 1 ] proposed the Berkeley Segmentation Dataset , a popular benchmark for segmentation and boundary detection that inspired many impactful works in perceptual grouping [ 2–5 ] . The recent surge of deep learning renders powerful representations with learned features using convolutional neural networks ( CNNs ) [ 6 ] . This has led to great advances in both areas [ 7–12 ] , but the two tasks are often considered separately . In light of the status quo , we consider coupled edge and segmentation learning . Our goal is twofold : ( 1 ) Multi-task learning - being able to produce high quality edge detection and segmentation . ( 2 ) Mutual improvement - the two tasks can help each other with non-trivial performance gains . Designing a principled framework is however nontrivial . The key question is how sparse edge signals can be effectively transformed into dense region-level ones to interact with segmentation . To this end , we propose a learnable recurrent message passing layer where semantic edges are considered ∗Equal contribution . Correspondence to Zhiding Yu < zhidingy @ nvidia.com > . †Work partially done during an internship at NVIDIA . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . as explicitly learned gating signals to refine segmentation . An overview of our framework is shown in Fig . 1 . Specifically , the dynamic message passing layer uses affinity gates to select the neighbor for message passing using max-pooling . It conducts message passing sweeps in each of the four directions : left to right , right to left , top to bottom and bottom to top . The message passing is jointly gated by both the affinity and edge gates , therefore allowing edge cues to naturally influence long-range dense predictions . As such , our framework presents a context module that is clean , compact yet powerful . Our technical contributions can be summarized as follows : • We formulate recurrent message passing as dynamic graph propagation ( DGP ) . We show that such a formulation simplifies the required normalization in propagation networks [ 13 ] . It also dynamically finds graph structures that encodes pixel affinities and improves dense prediction . • We propose a double-gate design where message passing is jointly gated by both the affinity and edge gates to refine the segmentation feature map . We show that this design together with the dynamic 1-way connection in DGP better couples segmentation and edge learning . • We obtain state-of-the-art results on joint semantic segmentation and edge detection . We also show that DGP leads to strong zero-shot robustness to natural corruptions with significant improvement over prior methods on Corrupted Cityscapes ( Cityscapes-C ) . Multitasking segmentation and edge learning is desirable for several reasons : 1 ) There are many downstream applications where both are needed , such as occlusion reasoning [ 14 ] , localization [ 15 ] , proposal generation [ 16–19 ] and conditional generation [ 20 ] . 2 ) There are many challenging cases where segmentation quality is poor but edge quality is far more superior , e.g . segmentation tends to be inferior near object boundaries since they are often optimized for IoU rather than precision [ 21 ] . In these cases , edge learning can potentially capture details missed by segmentation . 3 ) Implicitly improved model generalization as a result of the coupled learning [ 22 ] . 2 Related Work . Semantic Segmentation . There is a rich set of prior work in semantic segmentation . Long et al. , [ 7 ] proposed fully convolutional end-to-end training of semantic segmentation and pointed out its connection to recognition [ 6 ] . Chen et al . [ 8 ] introduced atrous convolution and atrous spatial pyramid pooling ( ASPP ) to capture multi-scale image contexts . Another contemporary work is the wide ResNet-38 which explores a relatively shallower but wider backbone [ 23 ] . It has also been shown that context plays an important role in segmentation , including context encoding [ 24 , 25 ] , multi-scale context [ 26–28 ] as well as relational context [ 29 , 30 ] . More recently , there has been a surge of interests in segmentation with Transformers [ 31–34 ] Boundary/Edge Detection . Similar to segmentation , boundary/edge detection have been widely studied as perceptual grouping problems in early literature [ 1–3 ] . Recent methods tend to resort to CNNs . For example , Bertasius et al . proposed a multi-scale deep network architecture for top-down contour detection , where as Xie et al . [ 11 ] further introduced holistically-nested edge detection ( HED ) for end-to-end edge learning . Besides detecting binary edges , Hariharan et al . [ 35 ] proposed the Semantic Boundaries Dataset ( SBD ) which has become a popular benchmark for semantic edge detection . Compared to binary edge detection , semantic edge detection involves the semantic classification of edge pixels in addition to localization which presents additional challenges to existing frameworks . A series of works including HFL [ 36 ] , CASENet [ 12 ] , SEAL [ 37 ] and STEAL [ 38 ] have followed up [ 35 ] and pushed the boundaries . Multi-task Segmentation and Edge Learning . Multi-task segmentation and edge learning remains under-studied but is not entirely new . Edges have been pairwise similarity cues to improve segmentation [ 3 , 39 , 18 ] and superpixels [ 4 ] . It was shown that edges can be transformed into dense regions through the Laplacian Eigenmaps of boundary transformed affinity [ 36 , 5 , 40 ] . Despite being robust , these methods are generally slow and not end-to-end trainable . For more recent CNN based models , end-to-end multi-task learning on a shared backbone is a natural choice [ 41 ] . In addition , proper regularization between segmentation and edge has been shown to improve the performance of both tasks . For example , Takikawa et al . [ 21 ] use softmax with temperature to impose consistency between segmentation and semantic edges . Zhen et al . [ 42 ] correlate semantic segmentation and edge detection tasks with a consistency loss . Our work goes beyond multi-tasking and let the two task more deeply coupled through dynamic graph propagation . Structured Dense Prediction . Segmentation is a dense prediction task where structured information can be useful . Structured prediction models such as Markov random fields ( MRFs ) [ 43 ] , conditional random fields ( CRFs ) [ 44–46 ] and energy minimization [ 47 , 48 ] have widely proved helpful in segmentation problems by imposing contrast-sensitive smoothness . In addition , structured inference can also be unfolded as network layers for end-to-end training [ 49 , 50 ] , therefore combining the advantages from both ends . Recently , there is an increasing trend to directly model the message passing process itself using multi-dimensional RNN [ 51–53 ] , graph RNN [ 54 , 55 ] and spatial propagation network ( SPN ) [ 13 , 56 , 57 ] . A common advantage is that these methods render more context-aware prediction with larger receptive fields while preserving local details similar to CRFs . Their ability to train end-to-end allows more powerful representation of inter-pixel similarities and thus better dense prediction quality . Our work can be broadly categorized into this category . It is worth mentioning that the contrast sensitive smoothness term in CRF and propagation networks [ 52 , 13 ] is an implicit modeling of edge signals by learning to relax smoothness constraints at high contrast areas . There have also been variants of SPN [ 58 ] that take binary edge as gate regularization . However , none of these works have explicitly addressed multi-tasking learning with category-aware edges whereas the proposed DGP framework presents a novel and effective solution to this problem . 3 Multi-task network . As a first step towards coupled segmentation and edge learning , we introduce two novel multi-task networks that are able to perform multi-task learning for both segmentation and edge detection . An overview of their architectures is shown in Fig . 2 . The rest of the paper follows these notations and covers more details for each module . Backbone . A CNN giving a semantic feature map ( denoted as F ) encoding the segmentation information . There could be multiple choices of the backbone networks , ranging from the popular architectures of DeepLabv2 [ 8 ] , DeepLabv3 [ 27 ] to latest state-of-the-arts such as DeepLabv3+ [ 28 ] and Dual-Attention Network [ 30 ] . Adopting powerful backbones will surely benefit the system-level performance , but the major purpose of this work does not completely lie in achieving state-of-the-arts . We are more interested in showing the effectiveness of a proposed framework on standard backbones . To this end , we consider both CASENet [ 12 ] and ResNet-38 [ 23 ] as the standard backbones for benchmarking . Both of them comprehensively covers the feature map resolutions of 1 , 1/2 , 1/4 and 1/8 , with the last resolution being a rule of thumb adopted by many segmentation networks . For both networks , we adopt atrous spatial pyramid pooling layers with dilation rate [ 6 , 12 , 18 , 24 , 30 , 36 ] to capture context information . We set the output channels to be 128 to produce a semantic feature map , fowllowed by either direct segmentation prediction ( baselines ) or the proposed DGP layer . Affinity stream . Convolution layers that aggregate the cross-layer information of the backbone and produces an affinity feature map ( denoted as A ) to encode the affinity of pairwise neighboring pixels . This stream is only used with the presence of recurrent message passing layer ( including its related baselines ) . The stream aims to model the simialrity of pair-wise neighboring pixels and serve as a major input to the gating signal in recurrent message passing layer . Compared to edge stream , the affinity stream seeks to capture coarser-level inter-pixel affinity with slightly deeper convolutions and higher dimensions . The resulting feature map A is a 256 dimensional tensor concatenated by ASPP and the side convolutional feature maps from the affinity stream . Edge stream . Dense skip-layer features with abundant details that are combined with edge classification layer through shared concatenation [ 12 ] to produce a semantic edge map ( denoted as E ) . We design an improved edge stream over CASENet [ 12 ] to better leverage detail information at multiple scales/levels with dense side connections . Unlike CASENet where the bottom blocks only provide 1-channel side features , we densely obtain side features from every sub-block residual modules . In CASENet , this gives side features with a total number of 31 dimensions . Similar rules apply to ResNet-38 , where the side feature has a total number of 17 dimensions We also found that applying sigmoid after side features with resolution 1/8 greatly benefits edge training from two aspects : ( 1 ) It helps to stabilize and prevent gradient explosions from dense skip connections . ( 2 ) It removes the typical aliasing edge issue caused by upsampling/deconvolution , and produces elegant smooth edges . Even though we do not explicitly apply techniques such as edge alignment [ 37 ] in this work , the proposed backbone is able to produce high quality edge predictions under noisy labels . We also notice that it is better to remove sigmoid for side features with higher resolutions . Edge prediction is multi-tasked alongside with ASPP using 1× 1 convolution on top of Res Block5 . This returns the K classes coarse edge predictions . Similar to CASENet , we upsample all side features together with the K-dimensional edge predictions to full image resolution and apply shared concatenation , where side features are repetitively shared and concatenated with each class for K times , followed by a K-way 3× 3 group convolution to produce semantic edges . Dynamic graph propagation . Learnable recurrent message passing layer that takes the above three branches as input to produce a refined semantic feature map ( denoted as F ∗ ) . More details regarding this module will be introduced in the next section .

Summary: This paper proposed a unified framework for detecting edges and segmenting objects. They conduct experiments on Cityscapes, SBD and Pascal VOC. The edge detection results are competitive to previous SOTA methods.