Learning 3D Dense Correspondence via Canonical Point Autoencoder

1 Introduction . With prior knowledge and experience , humans can easily perceive corresponding object parts ( e.g. , the wings from two different airplanes ) , understand their shape and appearance variance , in order to distinguish different objects coming from the same category . In computer vision , modeling dense correspondence between 3D shapes in one category is fundamental for numerous applications , such as robot grasping [ 1 , 2 ] , object manipulation [ 3 ] and texture mapping [ 2 , 4 ] . However , existing 3D cameras typically capture raw point clouds of shape surfaces that are arbitrarily-ordered and unstructured , in which correspondences are not established . 3D mesh representation , although is usually parameterized with UV maps that can indicate correspondences , can not be directly obtained from sensors and needs to be reconstructed from other types of representations , e.g. , 2D images [ 5 , 6 ] or 3D point clouds [ 7 ] . In this work , we focus on learning point cloud correspondences , which remains an open challenge since it is infeasible to label ground truth correspondence annotations . Without ground truth annotations , existing methods mainly discover shape correspondences via seeking a form of canonical space that can associate various instance shapes . For example , in particular shape domains such as human bodies [ 8 ] and human faces [ 9 ] , parameterized shape primitives have been designed to fit the observed raw data and to obtain the correspondences . Such designs , however , can not be generalized to other categories , e.g. , man-made objects [ 10 ] . Recently , several part co-segmentation networks relax the requirements of specific parameterized primitives , but instead decompose input shapes into an ordered group of simplest part constitutions [ 11 , 12 ] , in a self-supervised manner . These methods , however , require careful selection of the autoencoder architectures ( i.e. , they need to be considerably shallow to let the branches only able to represent simple shapes ) , and the number of part bases . Moreover , such part-based representation does not explicitly provide fine-grained ( e.g. , point-level ) correspondences . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this work , we introduce a novel canonical space where dense ( i.e. , point-level ) correspondences for all the shapes of a category can be explicitly obtained from . Inspired by 3D mesh representation [ 5 , 13 , 14 , 15 ] where shapes from one category are represented as deformations on top of a shape primitive , in our work , we set the canonical space as a 3D UV sphere . Our goal is to learn a “ point cloud-to-sphere mapping ” such that corresponding parts from different instances overlap when mapped onto the canonical sphere . In other words , similar to the mesh representation , a unique UV coordinate can represent the same semantic point/local region of shapes ( e.g. , the tip of an aeroplane ’ s wing ) , regardless of shape variations . Towards this goal , we introduce the canonical point autoencoder ( CPAE ) : we place the sphere primitive at the bottleneck ; the encoder non-linearly maps each individual input shape to the sphere primitive , where the decoder deforms the primitive back to match the original shape . We show that with several self-supervised objectives , this autoencoder architecture effectively ( 1 ) enforces the input points warped to the surface of the sphere primitives , and ( 2 ) encourages those corresponding points from different instances mapped to the same location on the sphere – both guarantee that the network learns correct dense correspondences . Essentially , we do not assume all object shapes in one category having the same topology , e.g. , an armchair does not have correspondence on its armrests , with another instance without an armrest . To introduce such uncertainty for correspondence matching , we propose an adaptive Chamfer loss on the bottleneck to allow customized primitive for each instance . As such , we are able to determine if a point on one instance has a correspondence in another point cloud . One advantage of the proposed method compared to the recent work [ 12 ] is that we can learn correspondences even when instances in the training dataset are not aligned , i.e. , our model is rotationinvariant within a certain rotation range and does not need to predict an additional rotation matrix as used in [ 12 ] . The main contributions of this work are : • We introduce a novel canonical space – a UV sphere , that explicitly represents dense correspondences of shapes from one category . • With the canonical space on the bottleneck , we design an autoencoder that learns such a “ point cloud-to-sphere mapping ” via a group of self-supervised objectives . • We apply the proposed method on various categories and quantitatively evaluate on the task of 3D semantic keypoint transfer and part segmentation label transfer , achieving comparable if not better performance than state-of-the-art methods . 2 Related Work . Deep Learning on Point Clouds . As a flexible and memory efficient representation of 3D shapes , point cloud has been widely studied and combined with deep neural networks . The PointNet [ 16 ] solves point cloud classification and segmentation by using MLPs and a max pooling layer to aggregate 3D shape information . One crucial property of the PointNet is that it is able to handle unordered input and is thus invariant to point permutations . In our work , we utilize the first few layers of the PointNet as our shape encoder . Another line of works [ 17 , 18 , 19 , 20 ] focus on reconstructing or generating point clouds . The most representative and related to our work is the FoldingNet [ 17 ] , where a point cloud is first encoded by a graph-based encoder and then reconstructed by sequentially applying the “ folding operation ” ( instantiated by MLPs ) to a 2D UV map . In this work , we reveal why the MLPs are able to preserve point order and utilize it as building blocks in our CPAE . Primitive Fitting . Primitive fitting has a long history in computer vision [ 21 ] . Several approaches focus on fitting individual samples with shape priors , such as blocks world [ 21 ] , generalized cylinders [ 22 ] , geons [ 23 ] , Lego pieces [ 24 ] , qualitative 3D blocks [ 25 ] , and patches [ 26 ] . However , these primitives are analyzed per-instance without order , therefore , can not be used for the correspondence . Some other approaches focus on fitting over a collection of shapes , which provide ordered set of primitives [ 11 , 27 , 28 , 29 , 30 , 31 , 32 ] . However , most of these approaches can only infer part-level correspondences and require additional supervision such as the number of primitives . In this work , we directly learns set-to-set alignment between point clouds , thus requires no shape prior . 3D Dense Correspondence . Given a pair of source and target instance in the same category , 3D dense correspondence learning targets at finding a corresponding point in the target instance for each point from the source instance . Several approaches [ 33 , 34 , 35 ] resolve this task through point cloud registration , using labeled pairwise correspondence as supervision . To relax constraints on supervision , Bhatnagar et al . [ 36 ] predict part correspondences to a template via implicit functions . Unfortunately , they require part labels for training . To unsupervisedly learn 3D dense correspondence , Chen et al . [ 37 ] propose a method to learn 3D structure points that are consistent across different instances . However , the model assumes structure similarity among different instances , which ignores intra-class variants and fails to detect non-existing correspondences between dissimilar shapes in the same category . Most relevant to our work , Liu et al . [ 12 ] introduce an unsupervised approach that leverages part features learned by the BAE-NET [ 11 ] to build dense correspondences . Their algorithm is able to calculate a confidence score representing the probability of correspondence . However , in order to train the implicit function , they require additional knowledge of object surface to compute the occupancy . Moreover , the correspondence learning in [ 12 ] heavily relies on the completeness of the part features from the BAE-NET [ 11 ] , which can lead to incorrect correspondences for parts that BAE-NET can not separate ( e.g. , flat surfaces , objects with fine-grained details ) . In contrast , our approach directly learns dense correspondences from the point cloud with self-supervision . Cycle Consistency . Our cross reconstruction loss is related to literature that uses cycle consistency as a supervisory signal to learn correspondence without ground truth annotations . Some of these works exploit the consistency among different modalities ( e.g. , 2D-3D ) . Specifically , [ 38 ] aim to learn cross-instance pairwise correspondence for 2D images using cycle-consistency terms guided by 3D synthetic data . [ 39 ] and [ 40 ] both leverage cycle-consistency terms to enforce correspondence between a 2D image and a 3D canonical template . Aside from these works that focus on crossmodality , work from [ 41 ] is mostly related to ours . They propose a “ point-wise ” cycle consistency loss that explicitly enforces a point deformed through any cycle of deformations to be mapped back to the origin location . Different from their approach , our consistency is instance-wise and we do not assume bijection mapping among shapes . 3D Deformable Mesh Representations . Our algorithm is also related to methods [ 5 , 13 , 14 , 15 ] that represent instance shapes as mesh deformations of a mesh primitive , i.e . a sphere template . Vertices on each instance surface that are mapped to the same locations on the shape primitive are discovered as correspondences . One obvious limitation is that only genus 0 shapes , e.g. , birds [ 5 , 13 ] can be deformed from a mesh sphere . In contrast , ours does not have such restriction thanks to the proposed non-linear mappings , instead of explicit deformation . 3 Proposed Method . In this section , we introduce our end-to-end CPAE for learning dense correspondences from point clouds without ground truth annotation . Given a point cloud SA with individual point p ∈ R3 ( see Figure 1 ) , our model : ( 1 ) predicts its canonical primitive UA ( a “ deformed ” point cloud with the same number of points as SA ) , which is supposed to be as close as possible to the canonical sphere in the bottleneck ; ( 2 ) reconstructs the original input point cloud back from UA . We show that while the first ensures each input point cloud to be warped to the surface of the sphere primitive , the second indirectly encourages the corresponding points from different point cloud instances to overlap during mapping to the primitive . In the following , we first describe our network architecture , i.e. , the encoder and decoder modules . We then introduce the adaptive Chamfer loss and the reconstruction losses that are applied to each individual module . Finally , we show that our decoder reconstructs ordered point clouds , i.e. , different point cloud instances can fetch their correspondences directly via the point indices , which provides a more accurate inference . 3.1 Network Architecture . Our method learns two mapping functions : one maps each individual 3D point p ∈ R3 in the world coordinates to the canonical space Φ ( p ) = u , u ∈ R3 , while the other conducts the inverse mapping Φ−1 ( u ) = p. Instead of using a reversible network , we instantiate Φ ( · ) and Φ−1 ( · ) with two MLPs respectively , and cascade them to construct an autoencoder ( i.e. , Φ ( · ) is the encoder and Ψ ( · ) ≈ Φ−1 ( · ) is the decoder ) . However , this mapping function can not be generalized to different point clouds if it only receives the coordinate of a single point p as the input . Therefore , we formulate them as conditional mapping functions by concatenating the input with a shape latent code ZA that represents a unique input shape SA , e.g. , Φ ( p , ZA ) , in order to generalize the mapping functions to all instances in the same category . There are three key modules in the proposed method : ( 1 ) a PointNet encoder that produces the shape latent code ; ( 2 ) an encoder MLP ( denoted as canonical mapping ) that maps a point cloud to a primitive sphere , and ( 3 ) a decoder MLP ( denoted as inverse mapping ) that deforms the primitive sphere back to the input point cloud . PointNet Encoder . As shown in Figure 1 , given a point cloud SA ∈ Rk×3 with k points , we encode it using a PointNet [ 16 ] encoder in a way similar to [ 8 , 18 , 19 ] . Each 3D point of the input point cloud is represented as a 512 dimensional vector using an MLP with 3 hidden layers of 64 , 128 , 512 neurons and ReLU activations . We then aggregate all point features with max-pooling and a linear layer to generate a global latent code zA ∈ R512 . We use PointNet because it produces a robust latent code representing the global shape , which is also invariant to input permutation . Canonical Mapping Φ ( · ) . To learn the canonical mapping function , we concatenate each point with the global latent code [ p , Z ] ∈ R515 as input to the MLP , which then outputs a 3D coordinate u ∈ R3 in the canonical space . We construct a sphere point cloud by uniformly sampling a large number of points from a standard sphere mesh , and place them as the canonical primitive at the bottleneck . A Chamfer loss is used to measure the difference between the outputted point u and primitive to encourage the mapped point adhering to the surface ( see Eq . ( 1 ) ) . Ideally , all the mapped instances from the same category should be as close as possible to the canonical sphere primitive . However , an instance may include parts/regions that do not exist in other instances due to intra-class variation , e.g. , not all chairs include armrests . With a conventional Chamfer loss , the shapes for objects with rare components do not converge during training , i.e. , those rare components are mapped to locations that are far away from the primitive surface . Thus , we relax the bidirectional constraint of the Chamfer loss and allow each instance to produce its own “ instance primitive ” to some extent ( see UA in Figure 1 ) . This is formulated via an adaptive Chamfer loss LACD : LACD ( UA , U ) = 1 |UA| ∑ p∈UA min q∈U ∥p− q∥2 + α 1 |U | ∑ q∈U min p∈UA ∥q − p∥2 ( 1 ) where α ∼ [ 0 , 1 ] is an adaptive parameter indicating to what extent the predicted instance primitive should match a canonical sphere , UA and U are the instance primitive and the canonical UV primitive ( e.g. , a 3D UV sphere ) , respectively . When α = 1 , LACD is equivalent to calculating the conventional Chamfer distance between the unfolded primitive and the canonical UV primitive . During training , α is initialized to 1 and gradually decreased to 0 . This allows the canonical mapping to predict instance-aware primitive since the second term of the Chamfer loss no longer constraints the primitive to be consistent with a canonical sphere . In the experiments ( see Section 4.5 ) , we show that this design allows us to infer rare object components that have no correspondences in other instances since non-corresponding regions usually occupy a distinct area in the canonical space . Inverse Mapping Ψ ( · ) . Similarly , we utilize another MLP in the decoder , which receives the concatenation of one point from the bottleneck and the global shape code , i.e. , [ u , Z ] ∈ R515 , as the input . The function learns to map it back to its original world coordinate p , which can be fulfilled via a point-to-point reconstruction loss . We leverage an MSE loss LMSE , a Chamfer distance LCD and an Earth Mover Distance ( EMD ) LEMD between the reconstruction Ŝ and input point cloud P : Lrec ( P , Ŝ ) = µ1LMSE ( P , Ŝ ) + µ2LCD ( P , Ŝ ) + µ3LEMD ( P , Ŝ ) ( 2 ) where µ1 , µ2 , µ3 are the weights , and empirically determined , µ1 = 1e3 , µ2 = 10 , and µ3 = 1 . The inverse mapping bears some resemblance to the principle of FoldingNet [ 17 ] , which adopts an MLP architecture to map one 2D coordinate together with a global shape representation to the 3D world1 . However , since FoldingNet does not have the forward mapping function as ours and the input points are untraceable , it needs to use Chamfer loss , instead of a point-wise loss ( Eq . ( 2 ) ) as ours . Align Canonical Mappings via CrossReconstruction . With the encoder and objective function , the canonical mapping is able to map the input points to the surface of the canonical primitive . However , even with the reconstruction loss ( Eq . ( 2 ) ) , there is no guarantee that the corresponding points from different shape instances can be overlapped on the sphere – which is the key to learn dense correspondences . For instance , in Figure 1 , while SA is mapped to the frontal half of the sphere , it is likely that another shape SB will be mapped to the rear half . To better align the mapped points , we introduce a cross-structured decoder that leverages : ( 1 ) a self-reconstruction branch as presented in Eq . ( 2 ) , and ( 2 ) a cross-reconstruction branch . As shown in Figure 2 , we feed the combination of the predicted instance primitive UA of point cloud SA , with the shape latent vector ZB of another randomly sampled point cloud SB to the folding decoder . We then minimize the Chamfer distance between the output shape ŜA→B and SB . The cross-reconstruction loss between the point cloud SA and SB is : Lcross ( SA , SB ) = LCD ( ŜA→B , SB ) ( 3 ) The cross-reconstruction branch ensures even by swapping the predicted primitives UA and UB , the inverse mapping can still reconstruct their own shapes , conditioned on their own shape latent codes . That is , the decoder encourages UA and UB to overlap on the canonical primitive , and thus ensures the network learns correct correspondences . Relation to Implicit Function . We note that since both mapping functions Φ ( · ) and Ψ ( · ) process each input point independently , they can be interpreted as conditional implicit functions , which are widely applied for 3D shape reconstruction [ 42 , 43 ] , view synthesis [ 44 , 45 ] , and recently for self-supervised 3D correspondence [ 12 ] . Such an interpretation reveals that the MLPs indeed learn continuous mappings that are feasible for interpolations , e.g. , once the decoder MLP is learned , it is able to map any point on a continues sphere surface that is not necessarily among the mapped input u , to the world coordinate . E.g , one can sample points densely from the surface of the sphere to reconstruct a continuous surface or a 3D mesh .

Summary: This paper presents a novel solution for learning dense correspondence between un-aligned 3D shapes in an self-supervised manner. The key contributions this paper is to design a novel canonical space for representing all 3D shapes and learning an auto-encoder to learn a mapping to and from an input point cloud to this canonical space. The effectiveness of this approach is shown for the 3D semantic keypoint transfer and part segmentation transfer tasks across different categories.