Shape-Tailored Deep Neural Networks Using PDEs for Segmentation

1 INTRODUCTION . Convolutional neural networks ( CNNs ) have been used extensively for segmentation problems in computer vision He et al . ( 2017 ) ; He et al . ( 2016 ) ; Chen et al . ( 2017 ) ; Xie & Tu ( 2015 ) . CNNs provide a framework for learning descriptors that are able to discriminate different textured or semantic regions within images . Much progress has been made in segmentation with CNNs but results are still far from human performance . Also , significant engineering must be performed to adapt CNNs to segmentation problems . A basic component in the architecture for segmentation problems involves labeling or grouping dense descriptors returned by a backbone CNN . A difficulty in grouping these descriptors arises , especially near the boundaries of segmentation regions , as CNN descriptors aggregate data from fixed shape ( square neighborhoods ) at each pixel and may thus aggregate data from different regions . This makes grouping these descriptors into a unique region difficult , which often results in errors in the grouping . In segmentation problems ( e.g. , semantic segmentation ) , current methods attempt to mitigate these errors by adding post-processing layers that aim to group simultaneously the ( coarse-scale ) descriptors from the CNN backbone and the fine-level pixel data . However , the errors introduced might not always be fixed . A more natural approach to avoid this problem is to consider the coarse and fine structure together , avoiding aggregation across boundaries , to prevent errors at the outset . To avoid such errors , one could design descriptors that aggregate data only within boundaries . To this end , Khan et al . ( 2015 ) introduced “ shape-tailored ” descriptors that aggregate data within a region of interest , and used these descriptors for segmentation . However , these descriptors are hand-crafted and do not perform on-par with learned approaches . Khan & Sundaramoorthi ( 2018 ) introduced learned shape-tailored descriptors by learning a neural network operating on the input channel dimension of input hand-crafted shape-tailored descriptors for segmentation . However , these networks , though deep in the channel dimension , did not filter data spatially within layers . Since an advantage of CNNs comes from exploiting spatial filtering at each depth of the network , in this work , we design shape-tailored networks that are deep and perform shape-tailored filtering in space at each layer using solutions of the Poisson PDE . This results in shape-tailored networks that provide more discriminative descriptors than a single shape-tailored kernel . This extension requires development of techniques to back-propagate through PDEs , which we derive in this work . Our contributions are specifically : 1 . We construct and show how to train ST-DNN , deep networks that perform shape-tailored spatial filtering via the Poisson PDE at each depth so as to generalize a CNN to arbitrarily shaped regions . 2 . We show analytically and empirically that ST-DNNs are covariant to translations and rotations as they inherit this property from the Poisson PDE . In segmentation , covariance ( a.k.a. , equivariance ) to translation and rotation is a desired property : if a segment in an image is found , then the corresponding segment should be found in the translated / rotated image ( or object ) . This property is not generally present with existing CNN-based segmentation methods even when trained with augmented translated and rotated images Azulay & Weiss ( 2019 ) , and requires special consideration . 3 . We show analytically and empirically that ST-DNNs are robust to domain deformations . These result from viewpoint change or object articulation , and so they should not affect the descriptor . 4 . To demonstrate ST-DNN and the properties above , we validate them on the task of segmentation , an important problem in low-level vision Malik & Perona ( 1990 ) ; Arbelaez et al . ( 2011b ) . Because of properties of the PDE , ST-DNN also have desirable generalization properties . This is because : a ) The robustness and covariance properties are built into our descriptors and do not need to be learned from data , b ) The PDE solutions , generalizations of Gabor-like filters Olshausen & Field ( 1996 ) ; Zador ( 2019 ) , have natural image structure inherent in their solutions and so this does not need to be learned from data , and c ) Our networks have fewer parameters compared to existing networks in segmentation . This is because the PDE solutions form a basis and only linear combinations of a few basis elements are needed to learn discriminative descriptors for segmentation . In contrast , CNNs spend a lot of parameters to learn this structure . 1.1 RELATED WORK . Traditional approaches to segmentation rely on hand-crafted features , e.g. , through a filter bank Haralick & Shapiro ( 1985 ) . These features are ambiguous near the boundaries of objects . In Khan et al . ( 2015 ) hand-crafted descriptors that aggregate data within object boundaries are constructed to avoid this , but lack sufficient capacity to capture the diversity of textures or be invariant to nuisances . Deep-learning based approaches have showed state-of-the-art results in edge-based methods Xie & Tu ( 2017 ) ; He et al . ( 2019 ) ; Deng et al . ( 2018 ) . Watershed is applied on edge-maps to obtain the segmentation . The main drawback of these methods is it is often difficult to form segmentations due to extraneous or faint edges , particularly when `` textons '' in textures are large . CNNs have been applied to compute descriptors for semantic segmentation , where pixels in an image are classified into certain semantic object classes Li et al . ( 2019 ) ; Huang et al . ( 2019 ) ; Du et al . ( 2019 ) ; Pang et al . ( 2019 ) ; Zhu et al . ( 2019 ) ; Liu et al . ( 2019 ) . Usually these classes are limited to a few object classes and do not tackle general textures , where the number of classes may be far greater , and thus such approaches are not directly applicable to texture segmentation . But semantic segmentation approaches may eventually benefit from our methodology as descriptors aggregating data only within objects or regions are also relevant to these problems . A learned shape-tailored descriptor Khan & Sundaramoorthi ( 2018 ) is constructed with a Siamese network on hand-crafted shape-tailored descriptors . However , Khan & Sundaramoorthi ( 2018 ) only does shape-tailored filtering in pre-processing as layering these requires new methods to train . We further examine covariance and robustness , not examined in Khan & Sundaramoorthi ( 2018 ) . Covariance to rotation in CNNs has been examined in recent works , e.g. , Weiler et al . ( 2018 ) ; Yin et al . ( 2019 ) ; Anderson et al . ( 2019 ) . They , however , are not shape-tailored so do not aggregate data only within shaped regions . Lack of robustness to deformation ( and translation ) in CNNs is examined in Azulay & Weiss ( 2019 ) and theoretically in Bietti & Mairal ( 2017 ) . Sifre & Mallat ( 2013 ) constructs deformation robust descriptors inspired by CNNs , but are hand-crafted . 2 CONSTRUCTION OF SHAPE-TAILORED DNN AND PROPERTIES . In this section , we design a deep neural network that outputs descriptors at each pixel within an arbitrary shaped region of interest and aggregates data only from within the region . We want the descriptors to be discriminative of different texture , yet robust to nuisances within the region ( e.g. , local photometric and geometric variability ) to be useful for segmentation . Our construction uses a Poisson PDE , which naturally smooths data only within a region of interest . Smoothing naturally yields robustness to geometric nuisances ( domain deformations ) . By taking linear combinations of derivatives of the output of the PDE , we can approximate the effect of general convolutional kernels but avoid mixing data across the boundary of region of interest . ST-DNN is also covariant to translations and rotations , inheriting it from the Poisson equation , which leads to the segmentation algorithm being covariant to such transformations . 2.1 SHAPE-TAILORED DNN DESCRIPTORS THROUGH POISSON PDE . Shape-tailored Smoothing via Poisson PDE : To construct a shape-tailored deep network , we first smooth the input to a layer using the Poisson PDE so as to aggregate data only within the region of interest , similar to what is done in Khan et al . ( 2015 ) for just the first layer . Let R ⊂ Ω ⊂ R2 be the region of interest , where Ω is the domain of the input image I : Ω→ Rk and k is the number of input channels to the layer . Let u : R → RM ( M is the number of output channels ) be the result of the smoothing ; the components u of u solve the PDE within R : { u ( x ) − α∆u ( x ) = I ( x ) x ∈ R ∇u ( x ) ·N = 0 x ∈ ∂R , ( 1 ) where I is a channel of I , ∂R is the boundary of R , N is normal to ∂R , α is the scale of smoothing and ∆/∇ are the Laplacian and the gradient respectively . It can be shown that the smoothing can be written in the form u ( x ) = ∫ R K ( x , y ) I ( y ) dy where K ( . , . ) is the Green ’ s function of the PDE , a smoothing kernel , which further shows that the PDE aggregates data only within R. Shape-tailored Deep Network : We can now generalize the operation of convolution tailored to the region of interest by taking linear combinations of partial derivatives of the output of the PDE equation 1 . This is motivated by the fact that in R = R2 , linear combinations of derivatives of Gaussians can approximate any kernel arbitrarily well . Gaussian filters are the solution of the heat equation , and the PDE equation 1 relates to the heat equation , i.e. , equation 1 is the steady state solution of a heat equation . Thus , linear combinations of derivatives of equation 1 generalize convolution to an arbitrary shape R ; in experiments , a few first order directional derivatives are sufficient for our segmentation tasks ( see Section 5 for details ) . A layer of the ST-DNN takes such linear combinations and rectifies it as follows : fi ( x ) = r ◦ Li ◦ T [ I ] ( x ) , ( 2 ) where I : R→ Rk is the input to the layer , T is an operator that outputs derivatives of the solution of the Poisson PDE equation 1 , Li ( y ) = wiy+bi is a point-wise linear function ( i.e. , a 1×1 convolution applied to combine different channels ) , r is the rectified linear function , and i indexes the layer of the network . Notice that since r and Li are pointwise operations , they preserve the property of T that it aggregates data only within the region R. We now compose layers to construct a ST-DNN as follows : F [ I ] ( x ) = s ◦ fm ◦ fm−1 ◦ fm−2 ◦ .... f0 ◦ I ( x ) , ( 3 ) where F [ I ] ( x ) is the output of the ST-DNN , f0 , ... , fm are the m+ 1 layers of the network , I is the input image , and s represents the soft-max operation ( to bound the output values ) . ST-DNN does not have a pooling layer because the PDE already aggregates data from a neighborhood by smoothing ; further , the lack of reduction in spatial dimension allows for more accurate shape estimation in our subsequent segmentation , and avoids the need for up-sampling layers . We will show that we can retain efficiency in training and inference .

Summary: The paper proposed a new "shape-tailored" convolutional layer for improving the accuracy of semantic segmentation. The shape-tailored layer is inspired by the Poisson partial differential equation which aggregate features from neighboring pixels through the linear combinations of partial derivatives of the output of PDEs. The paper proves several properties of the proposed ST-DNN and demonstrated good results in comparisons with previous semantic segmentation algorithms.