Contrastive Learning for Joint Normal Estimation
and Point Cloud Filtering

As mentioned earlier, contrastive learning has emerged as an effective method of generating latent representations of inputs such as images or point clouds based on similarity between augmented pairs of inputs which are seen by the network during training [Chen-SimCLR]. The work of Xie et al. [Xie-PointContrast] extended this method to 3D point clouds. Crucially, their work focuses on generating representations of entire point clouds, i.e., they use a global approach. However, as Guerrero et al. [Guerrero-PCPNet] point out, normal estimation at a given point relies on the local structure of the point neighborhood rather than the global structure of the entire point cloud. This is also true for the problem of point cloud filtering [Rakotosaona-PCN] as effective filtering mechanisms must preserve sharp feature information locally. This motivates our approach of developing a patch-based contrastive learning mechanism where noise corruption of input patches is used as an augmentation to develop different views of the same underlying clean patch. Thereafter, we employ the Normalized Temperature-scaled Cross Entropy (NT-XEnt) loss function detailed in Sec. 4.3 which promotes similarity of generated latent representations for a given positive pair of augmented patches. Inspired by the work of [Chen-SimCLR], we do not explicitly sample negative pairs as the remaining augmented pairs within a batch can be used for this purpose. Furthermore, the goal of this contrastive process is to bring representations of patches of the same underlying clean structure closer together, which is unlike a triplet based learning process which simultaneously brings representations closer for similar patches while pushing away representations of dissimilar patches.

Normal estimation and point cloud filtering are two interconnected tasks. Accurately predicted normals are central to reliable point cloud filtering and surface reconstruction as mentioned by [Lu-Deep-Feature-Preserving, RIMLS-Oztireli]. In a similar manner, predicting normals on less noisy patches provide more accurate results, as opposed to noisier patches where outliers affect the final prediction [Guerrero-PCPNet, Lenssen-Deep-Iterative]. This motivates our joint normal estimation and filtering approach where our regression network estimates patch normals along with point displacements to filter central patch points. Thereafter, the estimated normals are used to further refine the final filtered position. This approach helps exploit the interlinked relationship between normal estimation and point cloud filtering and motivates our joint approach.

Feature encoders trained using contrastive learning are adept at generating similar representations of similar inputs and dissimilar representations of dissimilar inputs. Guided by the intuition that two noisy variants of the same underlying clean patch should generate similar latent representations, we develop a contrastive learning framework with noise corruption as an augmentation to train a feature encoder $f(\ast)$ (see Fig. 1). This encoder is later used as part of a regression network (Fig. 3) to infer point normals and filtered displacements simultaneously. During training of the regression network, the pretrained feature encoder’s weights are kept frozen and only the regressor $h(\ast)$ is trained. The pretrained feature encoder is robust to noise and generates representations that can be consumed more effectively by the regression network during training. Fig. 2 illustrates t-SNE projections of 250 noisy variants of a given sharp feature patch and 250 noisy variants of a non-sharp feature patch obtained from the Cube shape in our dataset. The corresponding clean patches are illustrated at the top of Fig. 2. Each noisy variant contains Gaussian noise of standard deviation $\sigma$, ranging from $1.0\%$ to $2.5\%$ of the clean point cloud’s bounding box diagonal. We observe that a feature encoder trained without contrastive learning generates latent representations that are less similar for differing noisy patch variants sharing a common underlying clean structure, for both the sharp feature and non-sharp feature patch. This is evident from Fig. 2 as low noise patch variants (dark green markers) have projections that are far apart from their respective higher noise counterparts (light green/yellow markers).

The feature encoder pretrained using contrastive learning, with only noise corruption as an augmentation, generates latent representations whose t-SNE projections are clustered more closely, indicating that representations are more similar even as noise increases. This is due to the contrastive pretraining that exploits noise corruption as an augmentation and ensures robustness to noise. As such, latent representations generated by the feature encoder are easier for the regression network to distinguish, even for high noise patches, as these representations are similar to that of the underlying clean counterparts. Thereby, the contrastive learning based pretraining facilitates our joint normal estimation and point cloud filtering method.

We first introduce 3D patch-based contrastive learning to train a feature encoder capable of producing representations of point cloud patches (Fig. 1). This feature encoder consists of a PointNet-like architecture, with 5 Conv1D layers and a global max pool layer that generates a 1024 dimensional representation of input point cloud patches. These representations are projected to a 256 dimensional vector by a projection head consisting of 3 fully connected layers. Once the feature encoder has been trained, we use it to generate latent representations of input patches for regression tasks. Next, we train a regression network (Fig. 3) to predict patch normals and displacements simultaneously (i.e., normal and displacement of the central point of a patch). The regressor consists of the pretrained feature encoder and a MLP of 5 fully connected layers that output the desired normals and displacements. During the testing phase, these displacements are added to the initial point cloud, which produces a filtered point cloud that is refined and becomes the input for the next iteration of inference.

A clean point cloud consisting of $n$ points is described by $\mathbf{P}_{s}=\{p_{i}\ |\ p_{i}\in\mathbb{R}^{3},\ i\ =1,...,n\}$ where $s=1,\ldots,M$ enumerates all training shapes. A noisy point cloud $\hat{\mathbf{P}}_{s}$ can, thereafter, be characterized by the addition of noise onto the clean point cloud

where $\mathbf{N}(0,\sigma^{2})$ corresponds to additive Gaussian noise with a mean of $0$ and standard deviation of $\sigma$. For our training set, $\sigma$ takes on values of 0.25%, 0.5%, 1.0%, 1.5% and 2.5% of the bounding box diagonal length of $\mathbf{P}_{s}$. For a given shape, the set of 6 variant point clouds (1 clean and 5 noisy), is given by

Patches sampled from point clouds in $\mathbf{A}_{s}$ are utilized in the contrastive learning process.

An input patch centered at $p_{i}$, the $i^{th}$ point of a given point cloud $\mathbf{P}_{s}(\sigma_{1})\in\mathbf{A}_{s}$, can be described by

We create a contrastive pair $(\mathcal{P},\mathcal{Q})$ by randomly sampling another point cloud $\mathbf{P}_{s}(\sigma_{2})\in\mathbf{A}_{s}$ such that,

Here, $\mathbf{P}_{s}(\sigma_{1})$ and $\mathbf{P}_{s}(\sigma_{2})$ correspond to two noisy variants of the same underlying clean point cloud and $\sigma_{1}$ and $\sigma_{2}$ are chosen randomly. The process of pairing $\mathcal{P}$ with $\mathcal{Q}$, where both share the same underlying clean patch structure, as both are subsets of $\mathbf{A}_{s}$ and are centered at $p_{i}$, is the first augmentation. We note that $\mathcal{P}$ may have a different patch radius to $\mathcal{Q}$ as the bounding box diagonal length of $\mathbf{P}_{s}(\sigma_{1})$ differs from $\mathbf{P}_{s}(\sigma_{2})$ unless $\sigma_{1}=\sigma_{2}$. The patch radii $r_{\mathcal{P}}$ and $r_{\mathcal{Q}}$ are taken to be 5% of the respective bounding box diagonals. Patches are maintained at a fixed size to simplify the training procedure. Empirically, for each patch we sample 500 points. If the number of points is fewer, we increase it by copying random points in the patch and downsampling is applied otherwise. All patches are translated to the origin and normalized to [-1, 1], i.e., $\mathcal{P}=(\mathcal{P}-p_{i})/r_{\mathcal{P}}$ and $\mathcal{Q}=(\mathcal{Q}-p_{i})/r_{\mathcal{Q}}$.

For a batch of size $N$, we have a total of $2N$ augmented patches, i.e, we have input patches and their noise contrasted versions. The input patches are drawn from the training set $\mathbf{T}$, given by $\mathbf{T}=\bigcup_{s=1}^{M}\mathbf{A}_{s}$. All pairs $(\mathcal{P}_{k},\mathcal{Q}_{k})$ and $(\mathcal{Q}_{k},\mathcal{P}_{k})$, where $k=1,\ldots,N$, form positive contrastive pairs as these patches are drawn from point clouds of the same shape, the set $\mathbf{A}_{s}$, and correspond to the same central point $p_{i}$. Any pairs $(\mathcal{P}_{k},\mathcal{Q}_{l})$ or $(\mathcal{Q}_{l},\mathcal{P}_{k})$ where $k,l=1,\ldots,N\land k\neq l$ constitute negative pairs. For each augmented patch, there exists $1$ positive pair and $2(N-1)$ negative pairs. We do not pair the augmented patch with itself, i.e., we do not consider $(\mathcal{P}_{k},\mathcal{P}_{k})$ or $(\mathcal{Q}_{k},\mathcal{Q}_{k})$.

Positive contrastive pairs are put into the canonical basis by taking the inverse of the eigenvector matrix, obtained from the covariance matrix of $\mathcal{P}_{k}$, Eq. (5), and matrix multiplying both patches by it.

Finally, we apply a second augmentation to $\mathcal{Q}_{k}$ by randomly rotating it around either the $x$, $y$, or $z$ axis by an angle $\theta\in\{0,\pi/12,\pi/6,\pi/4,\pi/3,\pi/2,7\pi/12,2\pi/3,3\pi/4,5\pi/6,\pi\}$. This second augmentation allows the network to learn patch similarity based on patch structure despite rotations and reinforces the effect of contrastive learning.

Once contrastive pairs have been generated, we train our feature encoder $f(\ast)$ in a contrastive learning manner. As patches within a positive pair correspond to the same ground-truth patch, albeit with different levels of additive noise and arbitrary rotations applied to them, they should produce similar representations. As such, we adopt the Normalized Temperature-scaled Cross Entropy (NT-XEnt) loss [Chen-SimCLR] to achieve this. The loss for a positive pair $(\mathcal{P}_{k},\mathcal{Q}_{k})$ in the batch is given by

where $l=1,\ldots,2N$ enumerates over all projections. The function $\mathbbm{1}_{[k\neq l]}=0$ if $k=l$ and $1$ otherwise. We empirically set $\tau=0.01$ based on results. In Eq. (6), the projection $\bm{z}_{\ast}=g(f(\ast))$ with $g(\ast)$ and $f(\ast)$ parameterized by the projection head and feature encoder, respectively. Finally, the contrastive loss for the entire batch is,

where $k=1,\ldots,N$ enumerates over all positive pairs. Fig. 1 shows how a single positive pair within a batch is processed by the feature encoder and projection head. All positive and negative pairs within the batch see the same copy of the feature encoder and projection head and all pairs are processed simultaneously. The larger the batch size, the greater the number of negative pairs available for the contrastive loss calculation. Once the training of the contrastive learning network has been completed, we discard the projection head and use the feature encoder as the backbone for our regression network. Representations generated by the feature encoder are consumed by the regressor, comprising 5 MLPs, and outputs a 2-tuple of displacement and normal vectors. During the training of the regressor, the weights for the feature encoder are kept frozen. The regression process is visualized in Fig. 3.

We propose a novel, joint position and normal based loss for training our regression network. This joint loss allows our network to, when trained, perform both a position update of a noisy point while also estimating its corresponding normal. Therefore, our regressor loss function comprises two main components: $L_{pos}$, the position loss and $L_{normal}$, the normal loss.

Approximating the clean surface. The position loss term is inspired by the work of [Rakotosaona-PCN]. We consider all points within the ground-truth patch $\mathcal{P}^{\star}$ and seek to minimize the squared $L_{2}$ norm instead of solely using a single fixed target.

where $\tilde{p}_{i}$ is the filtered point from the regressor given a noisy patch $\mathcal{P}$ centered at $p_{i}$, i.e., the $i^{th}$ point in $\mathcal{P}$. The ground truth patch $\mathcal{P}^{\star}=\{p_{j}\ |\ p_{j}\in\mathbf{P}_{s}\land\left\|p_{j}-p_{i}\right\|_{2}<r_{\mathcal{P}}\}$ is translated to the origin and normalized such that $\mathcal{P}^{\star}=(\mathcal{P}^{\star}-p_{i})/r_{\mathcal{P}}$. This loss term allows the regressor to filter the noisy central point back to the clean patch surface. However, this does not ensure that the filtered point is centered within the ground truth patch which leads to unwanted clustering of filtered points.

Ensuring regular distribution of points. To avoid unwanted point clustering, we employ a regularization term which promotes the centering of filtered points within ground truth patches. By ensuring that filtered points lie close to the central points of their corresponding ground truth patches, we recover a regular distribution of points within the filtered point cloud.

Intuitively, if the filtered point lies away from the ground truth patch center, this leads to a larger penalization due to Eq. (9). For example, given a circular patch with radius $r$, the minimum value of Eq. (9) is $r^{2}$ where the inferred point lies at the center of the ground truth patch and the furthest point is $r$ away. $L^{1}_{pos}$ and $L^{2}_{pos}$ form the position based loss contribution to the final loss

where $\beta$ is the parameter that controls the regularization term’s contribution to the position loss. We empirically set it to $0.01$, as we notice that a large contribution affects the convergence of the final loss function.

Normal estimation. We then develop a relationship between a regressed position and its normal. Intuitively, if the ground-truth patch point, which minimizes the squared $L_{2}$ norm between positions, corresponds to the true position of the filtered point, then that point’s normal should correspond to the true normal of the filtered point. If the ground-truth patch’s central point which minimizes the squared $L_{2}$ distance is given by

then the angle difference between the predicted normal and the ground-truth normal can be expressed in terms of cosine similarity between the two:

where $n_{\tilde{p}_{i}}$ and $n_{p^{*}_{j}}$ correspond to the normals at $\tilde{p}_{i}$ and $p^{*}_{j}$, respectively. This cosine term can now be used to construct our normal loss:

The loss function in Eq. (13) is a periodic function which penalizes $\theta$ values away from $0$ and $\pi$. Therefore, it encourages the predicted normal to be as close to the ground-truth normal as possible. It also assumes that the predicted normal with an angle difference of $\pi$ is equivalent to the ground-truth normal. The $\delta$ term, which is empirically set to $0.3$, serves to control the shape of loss function $L_{normal}$, wherein, angle differences close to $\pi/2$ are heavily penalized. This penalization decreases closer to $0$ and $\pi$. Finally, we express our joint loss as

where $\alpha$ controls the relative contributions of the position and normal losses to the final loss function. We empirically set $\alpha$ to $0.9$ as the emphasis for the regressor is denoising the point cloud iteratively. This in turn leads to better normal estimation results.

We also examine Eq. (17), a variant of the joint loss function, Eq. (14), which utilizes the point-to-point correspondences between ground truth points $p^{\star}_{i}$ and filtered points $\tilde{p}_{i}$, i.e., using fixed ground truth targets.

where $\cos(\theta)=n_{\tilde{p}_{i}}\cdot n_{p^{\star}_{i}}$ with $n_{\tilde{p}_{i}}$ and $n_{p^{\star}_{i}}$ being the predicted normal and ground truth normal, respectively. As we regress the filtered point directly back to the ground truth central point, $L^{alt}_{pos}$ does not contain a repulsion term. The regressor trained with this loss function performs sub-optimally to that of the regressor trained using Eq. (14). As noted in [Rakotosaona-PCN], multiple clean points, within a given neighbourhood, may be perturbed in such a way as to result in the same noisy point. Therefore, the $L_{2}$ norm minimization between a filtered point and ground truth point, Eq. (15), cannot successfully remove the noise component tangential to the surface and leads to a lower filtering performance. This motivates our use of Eq. (14) which regresses the filtered point back to the surface while ensuring it is centered as best possible within the ground truth patch. More details are given in Sec. LABEL:sec:alt-joint-loss.

The regression network $h(\ast)$, trained based on the loss function defined by Eq. (14), outputs a 6D vector or 2-tuple $(h_{0},h_{1})$ of 3D vectors. The first element corresponds to the displacement required to obtain the filtered point $p^{\prime}_{i}$, and the second to the corresponding normal vector $\tilde{n}_{i}$. As these two vectors are in the space defined by the eigenvectors of the patch covariance matrix, they must first be transformed back to the original space. Subsequently, the filtered point is given by $p^{\prime}_{i}=p_{i}+T^{-1}(h(f(T(\mathcal{P})))_{0})\cdot r_{\mathcal{P}}$ with the original noisy point $p_{i}$. The normal vector in the original space is $n_{\bar{\bar{p}}_{i}}=T^{-1}(h(f(T(\mathcal{P}))))_{1})$ where $f(\ast)$ is the feature encoder.

Given $p^{\prime}_{i}$, we apply refinement during post-processing to obtain the final filtered point as suggested by [Rakotosaona-PCN, Lu-Low-Rank]. The first is a Taubin smoothing-like inflation step [Rakotosaona-PCN] and the second is the LRMA position update [Lu-Low-Rank] to combat shrinking of the point cloud and avoid incorrect displacement of points along the surface, respectively. The ablation study on the post-processing refinement is discussed in Section LABEL:sec:no-post-processing. The new position after applying the inflation step is given by,

where we take $p^{\prime}_{j}\in\mathcal{N}(p^{\prime}_{i})$ as the neighborhood of 100 filtered points in the vicinity of $p^{\prime}_{i}$ and $p_{j}$ is the original point position before filtering. The LRMA position update yields our final filtered point. The position update is given by,

with $\mathcal{N}(\bar{p}_{i})$ being the neighborhood of 20 points in the vicinity of $\bar{p}_{i}$.

Our training set consists of 22 synthetic point clouds (Fig. 4): 11 CAD shapes and 11 non-CAD shapes. The validation set consists of 3 shapes, 1 CAD and 2 non-CAD shapes, while the test set consists of 23 shapes: 14 CAD and 9 non-CAD. Each shape is a point cloud of 100K points, which have been randomly sampled from their original surfaces. For training, we create 5 additional noisy variants of each training shape by adding Gaussian noise with standard deviations of 0.25%, 0.5%, 1.0%, 1.5% and 2.5% of the clean point cloud’s bounding box diagonal. These 6 variants (clean and 5 noise levels) for each shape give a total of 132 point clouds for training purposes. For validation, we consider 2 noisy variants of the initial 3 clean validation shapes, with 0.5% and 1.0% noise, resulting in a total of 6 validation point clouds. In the testing phase, we examine the robustness of our model at unseen noise levels by utilizing 0.6%, 0.8%, 1.1%, 1.5% and 2.0% Gaussian noise for each test shape, yielding 115 test point clouds.

Sharp features. To evaluate performance at sharp features, we classify points within our synthetic dataset as feature and non-feature points. Please refer to the supplementary document for more details.

The contrastive learning and regression networks are both trained on NVIDIA A100 GPUs using PyTorch 1.7.1 with CUDA 11.0. The contrastive learning network is trained for 150 epochs, with the Adam optimizer and a learning rate of $3\times 10^{-4}$. The regression network is trained for 30 epochs, utilizing the SGD optimizer with a learning rate of $1\times 10^{-2}$.

We compare our method with state-of-the-art normal estimation and point cloud filtering methods. For normal estimation, we consider conventional PCA [Hoppe-PCA], HoughCNN [Boulch-HoughCNN], NINormal [Wang-NINormal], PCPNet [Guerrero-PCPNet], Nesti-Net [Ben-Shabat-Nesti-Net], Deep Iterative (DI) [Lenssen-Deep-Iterative] and AdaFit [Zhu-AdaFit] on our test set including synthetic and scanned point clouds. We do not compare with DeepFit [Ben-Shabat-DeepFit] as AdaFit is based on DeepFit and achieves better results. PCA requires the manual selection of neighborhood sizes for plane-fitting. For synthetic shapes, we utilize three different neighborhood sizes, i.e., 60 points for 0.6% Gaussian noise, 150 points for 0.8% Gaussian noise and 200 points for 1.1%, 1.5% and 2.0% Gaussian noise. For scanned surfaces a neighborhood of 100 points is used. Our metric for comparison is the Mean Squared Angular Error (MSAE) where we calculate angle differences between ground truth normals $n_{i}$ and predicted normals $\tilde{n}_{i}$ and take the mean of their squares.