Point Cloud Instance Segmentation using Probabilistic Embeddings

Introduction

In this paper we tackle the problem of instance segmentation of point clouds. In instance segmentation we would like to assign two labels to each point in a point cloud. The first label is the class label (\eg, leg, back, seat, … for a chair data set) and the second label is the instance ID (a unique number, \eg, to distinguish the different legs of a chair). While instance segmentation had many recent successes in the image domain , we believe that the problem of instance segmentation for point clouds is not sufficiently explored.

We build our work on the idea of embedding-based instance segmentation, that is very popular in the image and volume domain and has also been successfully applied in the point clouds domain . In this approach typically two steps are employed. In the first step, each point (or pixel) is embedded in a feature space such that points belonging to the same instance should be close and points belonging to different instances should be further apart from each other. In the second step points are grouped using a clustering algorithm, such as mean-shift or greedy clustering.

One important design choice in embedding-based methods is the dimensionality of the feature space. Some methods propose to use a high dimensional feature space , while others use a low dimensional features space that has the same dimensionality as the input data , \eg, 2D for images, and 3D for point clouds. Methods with a low dimensional embedding space not only have lower computational complexity, but they also lead to better interpretability, \eg, embeddings are encoded as offset vectors towards instance centers.

Therefore, the main goal of our work is to extend the expressiveness of the embedding space in a way that leads to improved segmentation performance. Our proposed solution is to employ probabilistic embeddings, such that each point in the embedding space is encoded by a distribution. While assessing uncertainty is a popular tool in recent computer vision research and we introduce this idea to the task of instance segmentation. Incorporating uncertainty leads to an important improvement in segmentation performance. For example, on the PartNet fine-grained instance segmentation dataset we can improve the SOTA by 3.1% average per-category mAP.

In the remainder of the paper, we will give more details on the probabilistic embedding algorithm (Sec. 3), explain the embedding step (Sec. 3.1) and the clustering step (Sec. 3.4) in more detail.

We propose to use probabilistic embeddings for instance segmentation and present a complete framework in the context of point cloud instance segmentation based on probabilistic embeddings.

We develop a new loss function for the clustering step that is especially suited for high-granularity data sets.

We show that the proposed probabilistic embeddings can be incorporated into existing embedding-based methods.

Related work

The dominant approaches for image instance segmentation are proposal-based methods , which are built upon object detection methods . Typically, they have higher quality, but a slower computation time compared to proposal free methods. The mainstream proposal free approaches are based on metric learning. The basic idea is to learn an embedding space, in which pixels belonging to the same object instance are close to each other and distant to pixels belonging to other object instances . All above works are based on high-dimensional embedding, while more recent works show that 2D spatial embedding is sufficient to achieve the same or even higher performance.

2 3D point cloud instance segmentation

SGPN uses PointNet++ as backbone network and designs a double-hinge loss function to learn a pairwise similarity matrix of points. GSPN produces object proposals with high objectness for point cloud instance segmentation. ASIS is a module capable of making semantic segmentation and instance segmentation take advantage of each other. release a large scale point cloud dataset for part instance segmentation and benchmark their method and SGPN on this dataset. PointGroup and OccuSeg achieved great success in scene datasets by voxelizing point clouds.

3 Uncertainty in computer vision

present a unified framework combining model uncertainty with data uncertainty and can estimate uncertainty in classification and regression tasks. We introduce uncertainty estimation to the literature of instance embedding, by modeling points as random variables. Our method is related to recent works in deep generative networks . They use a stochastic encoder to encode a data sample as a set of random variables, while focusing on solving the problem of backpropagation through random variables in deep neural networks. We deal with this problem by using a probabilistic product kernel .

Method

A training sample is a labeled 3D point cloud. It consists of point coordinates $\left\{\mathbf{x}_{i}\right\}^{N}_{i=1}$, class labels $\left\{y_{i}\right\}^{N}_{i=1}$ and instance IDs $\left\{z_{i}\right\}^{N}_{i=1}$. We want to train a neural network to infer per point class labels and per point instance IDs at the same time.

A common approach in the literature of instance segmentation is to learn a function to embed pixels/points into a space where pair-wise similarity can be measured. Usually, this function is a deep neural network $f$ which transforms an unordered point set $\left\{\mathbf{x}_{i}\right\}^{N}_{i=1}$ to embeddings $\left\{\mathbf{e}_{i}\right\}^{N}_{i=1}$.

Instead of deterministic embeddings used in previous work, here we consider a probabilistic embedding, by modeling $\mathbf{e}_{i}$ as a random variable, $\mathbf{e}_{i}\sim p_{i}(\mathbf{e}),$

where $p_{i}$ is a probability density function. In Section 3.3 we will need to calculate the sum of random variables. In the ideal case, the distribution of a single random variable and the sum of multiple random variables has the same type of distribution that can be described with a few parameters. We choose to work with the tri-variate Gaussian distributionRefer to for a discussion why spatial embedding works. $p_{i}(\mathbf{e})=\mathcal{N}(\mathbf{e};\boldsymbol{\mu}_{i},\boldsymbol{\Sigma}_{i})$

2 Similarity measure

In deterministic embeddings, the (dis)similarity between points is usually measured by Euclidean distance $\left\|\mathbf{e}_{i}-\mathbf{e}_{j}\right\|,$ or cosine similarity $\frac{\mathbf{e}_{i}^{\intercal}\mathbf{e}_{j}}{\left\|\mathbf{e}_{i}\right\|\left\|\mathbf{e}_{j}\right\|}$ (See Figure 2). Since now we are using probabilistic embeddings, a similarity kernel for random variables needs to be selected. Here we describe the Bhattacharyya kernel .

We choose this kernel as our similarity measure for two reasons, 1) the Bhattacharyya kernel is symmetric, i.e. $\mathcal{K}(p,q)=\mathcal{K}(q,p)$; 2) the Bhattacharyya kernel has values between (no similarity) and $1$ (maximal similarity). And $\mathcal{K}(p,q)=1$ if and only if $p=q$.

Then the similarity $\kappa(\cdot,\cdot)$ between random variables can be represented by the Bhattacharyya kernel of their probability density functionsRefer to for a derivation.,

If the uncertainties $\boldsymbol{\sigma}_{i}$ and $\boldsymbol{\sigma}_{j}$ have a large difference, $\beta_{i,j}$ will be small, so will be $\kappa(\mathbf{e}_{i},\mathbf{e}_{j})$. See Fig. 4.

If the centers $\boldsymbol{\mu}_{i}$ and $\boldsymbol{\mu}_{j}$ have a large difference, the exponential term will be small, so will be $\kappa(\mathbf{e}_{i},\mathbf{e}_{j})$. See Fig. 4.

The scale term $\beta_{i,j}=1$ if and only if the uncertainties $\boldsymbol{\sigma}_{i}$ and $\boldsymbol{\sigma}_{j}$ are element-wise equal. In this case, $\kappa(\mathbf{e}_{i},\mathbf{e}_{j})$ becomes an anisotropic Gaussian kernel,

The exponential term equals $1$ if and only if the centers $\boldsymbol{\mu}_{i}$ and $\boldsymbol{\mu}_{j}$ are element-wise equal. In this case, $\kappa(\mathbf{e}_{i},\mathbf{e}_{j})$ becomes $\beta_{i,j}$, i.e., the similarity between uncertainties. This property allows two points that have the same embedding centers to have a low similarity, as long as $\beta_{i,j}$ is small.

In the following, we discuss multiple choices of embedding distributions that we will evaluate in Sec 4.2.

The embeddings $\{\mathbf{e}_{i}\}^{N}_{i=1}$ are homoscedastic if they have the same variance $\boldsymbol{\Sigma}$. In this case, for a point cloud $\mathbf{X}$ we learn to predict a single $\boldsymbol{\Sigma}$ instead of point-dependent variances $\left\{\boldsymbol{\Sigma}_{i}\right\}^{N}_{i=1}$. And the similarity kernel becomes $\kappa(\mathbf{e}_{i},\mathbf{e}_{j})=\exp\left(-\left\|\boldsymbol{\mu}_{i}-\boldsymbol{\mu}_{j}\right\|^{2}_{\boldsymbol{\Sigma}_{i,j}^{-1}}\right)$, which is also the form of the RBF kernel in Eq. (2).

Isotropy vs. Anisotropy.

The variance $\boldsymbol{\Sigma}_{i}$ is isotropic if its diagonal elements (variances of dimensions) are the same. Then we can write $\boldsymbol{\Sigma}_{i}=\sigma_{i}^{2}\mathbf{I}$, where $\mathbf{I}$ is a $3\times 3$ identity matrix. The similarity can be written as

where $\beta_{i,j}=\left(\left(\sigma_{i}/\sigma_{j}+\sigma_{j}/\sigma_{i}\right)/2\right)^{-\frac{3}{2}}$ and $\alpha_{i,j}=4(\sigma_{i}^{2}+\sigma_{j}^{2})$.

3 Instance grouping

Let $\left\{i:z_{i}=k\right\}$ be the index set of points having instance ID $k$. We take an average of these embeddings to get the embedding $\mathbf{c}_{k}$ of instance $k$, $\mathbf{c}_{k}=\frac{1}{\left|\{i:z_{i}=k\}\right|}\sum_{\{i:z_{i}=k\}}\mathbf{e}_{i}$. Since the sum of Gaussian random variables is still a Gaussian random variable, we can derive the following:

Now we can measure the similarity between a point and an instance by using $\kappa(\mathbf{e}_{i},\mathbf{c}_{k})$.

If $z_{i}=k$, we want $\kappa(\mathbf{e}_{i},\mathbf{c}_{k})$ to be close to 1, otherwise 0. We can optimize a binary cross entropy loss function,

However, in practice, this suffers from a serious foreground-background imbalance problem. To remedy this drawback we propose to use the combined log-Dice loss function instead:

where $\mathds{1}_{z_{i}=k}$ is an indicator function which equals $1$ when $z_{i}=k$, otherwise.

As we can see in Figure 4, when $\sigma_{i}^{(l)}$ and $\sigma_{j}^{(l)}$ goes to infinity while keeping $\sigma_{i}^{(l)}=\sigma_{j}^{(l)}$, $\beta_{i,j}=1$ and the similarity equals to 1 no matter what the value $\boldsymbol{\mu}_{i}-\boldsymbol{\mu}_{j}$ is. Formally speaking,

Consequently, the similarity degenerates to constant $1$ for every pair of embeddings. To address this issue, we propose an entropy regularizer,

4 Semantic classification

Neven et al. introduces a way to use score maps to find cluster centers. Our main novelty is the new loss function, so our description focuses on this part. We still describe the greedy clustering steps from for completeness. In Section 4.2, we compare our new center-aware loss to the previously used MSE loss in .

$\mathbf{P}[i,:]$ is a probability vector and can be used to infer class label $y_{i}$ of point $\mathbf{x}_{i}$, i.e., $y_{i}=\arg\max_{l=1}^{L}\mathbf{P}[i,l]$.

For foreground class labels $l\in\{1,2,\dots,L\}$, $\mathbf{P}[:,l]$ is a score map of being an instance center with class label $l$.

The first condition is easy to satisfy with the cross entropy loss. Assuming $\mathbf{P}[i,:]$ is the output of a softmax function, we can minimize,

where $\mathds{1}_{y_{i}=l}$ is an indicator function which equals $1$ when $y_{i}=l$, otherwise.

where $y(k)$ is the class label of instance $k$, due to the fact that $\left\{\mathbf{x}_{i}:z_{i}=k\right\}$ must have the same class label. (See an illustration in Figure 5.) After that, we want both $\mathbf{P}[:,l]$ and $\mathbf{Q}[:,l]$ to achieve local maxima at the same points for all $l\in\{1,2,\dots,L\}$. When we are doing inference, these local maxima are chosen as instance centers. Therefore, the first condition can be weakened, and only points which are close to instance centers should be classified correctly.

We design a new loss function to satisfy the two conditions at the same time,

Here $\mathbf{Q}$ is fixed as a target when training. We can view $\mathcal{L}$ in two ways,

First, we switch the order of summation in Eq. (12),

The value of this quantity $-\mathbf{Q}[i,l]\log\mathbf{P}[i,l]$ is high when weight term $\mathbf{Q}[i,l]$ is high, and if we minimize it, we are forcing $-\log\mathbf{P}[i,l]$ to be small. Consequently, $\mathbf{P}[i,l]$ would be large. This guarantees local maxima of $\mathbf{Q}[:,l]$ are also local maxima of $\mathbf{P}[:,l]$. And minimizing this loss term is equivalent to minimize the KL-divergence between (unnormalized probability) $\mathbf{Q}[:,l]$ and (unnormalized probability) $\mathbf{P}[:,l]$,

Second, we look at the inner summation of Eq. (12),

The inner summation inside the round bracket is the cross entropy between $\mathbf{Q}[i,:]$ and $\mathbf{P}[i,:]$. And it is equivalent to replacing the one-hot vector in Equation 9 with $\mathbf{Q}[i,:]$. Also, it is the form of label smoothing, a commonly used training trick in image classification . The closer $\mathbf{Q}[i,:]$ is to a one-hot vector, the more confidence we give to the classification loss of point $\mathbf{x}_{i}$. By definition of $\mathbf{Q}[i,l]$, it can be easily seen that the resulting classifier only classifies near-centers points correctly. Thus we call our new loss function the center-aware loss.

The inference process is done with a greedy approach . From foreground score maps $\left\{\mathbf{P}[:,1],\mathbf{P}[:,2],\dots,\mathbf{P}[:,L]\right\}$, we sample a point $\mathbf{x}_{i_{0}}$ with highest score $\mathbf{P}[i_{0},l_{0}]$, where $i_{0}$ is the point index and $l_{0}$ is its class label,

The point $\mathbf{x}_{i_{0}}$ is an anchor and we want to find all similar points. Specifically we find all points $\mathbf{x}_{i}$ with $\kappa(\mathbf{e}_{i},\mathbf{e}_{i_{0}})\geq\tau.$ As a result, the instance ID of $\mathbf{x}_{i}$ is . After that, all points satisfying the inequality are all masked out. Similarly, we sample $\mathbf{x}_{i_{1}}$ and mask out points with instance ID $1$, sample $\mathbf{x}_{i_{2}}$ and mask out points with instance ID $2$, and so on. We stop this loop if there is no point left. We use the validation set to fit hyperparameter $\tau$, which is $0.35$ in our experiments.

5 Implementation

For a fair comparison to our main competitor PartNet we keep as much of their structure as possible (Note that PartNet is the name of a dataset as well as an instance segmentation method). We also use PointNet++ as the feature extraction backbone, with the same parameters as . We use 3 output heads for centers, uncertainties, and scores as in $f(\left\{\mathbf{x}_{i}\right\}^{N}_{i=1})=\left\{\boldsymbol{\mu}_{i},\boldsymbol{\sigma}_{i},\mathbf{p}_{i}\right\}^{N}_{i=1}$. We list the activation functions for output heads in Table 2.

We use random jittering, translation (between -0.01 and 0.01) and rotation (between $-15^{\circ}$ and $15^{\circ}$ for each axis) as data augmentation, and use the Adam optimizer. We use a batch-size of $16$ and an initial learning rate of $0.001$ for $500$ epochs with a decay factor of $0.5$ at epoch 50 and epoch 150.

Results

PartNet provides coarse-, middle- and fine-grained part instance-level annotations for 3D point clouds from ShapeNet . It contains 24 object categories, but the number of training samples varies greatly from 92 to 5707 for different categories. In contrast to indoor scene point cloud datasets (\eg, ScanNet by ), instances (object parts) of PartNet require more context to be classified and are connected. Many visually alike parts have different semantic labels, \eg, ping-pong table’s legs and pool table’s legs in the category of table. Also, instance masks should have no overlaps. All these make it a very challenging dataset for instance segmentation.

We report per-category mean Average Precision (mAP) scores for the PartNet dataset in Table 1. The IoU threshold is $0.5$. We compare our probabilistic embedding algorithm to PartNet and SGPN . The results are averaged over three levels of granularity (fine(3), middle(2), and coarse(1)).

On the complete dataset, our method outperforms the best competitor PartNet by 3.1% average per-category mAP. We can observe that our method has a slightly bigger advantage in fine-grained instance segmentation compared to coarse-grained instance segmentation (3.2% vs. 2.5%). We can also observe consistent improvements in categories with little as well as many training samples. While we beat SOTA in all categories with many training samples (Chair, Table, StorageFurniture, and Lamp), PartNet has better results in some of the categories with fewer training samples.

We also show visualization examples in Figure 6. Compared to PartNet , our method shows great improvement especially when there are many instances in a point cloud.

ScanNet.

As baseline method we chose a network based on performance and availability of code. Since the best methods, such as OccuSeg, do not release code for ScanNet, we decided to build our own baseline using MinkowskiNet as feature backbone. MinkowskiNet is a sparse tensor network that achieved great results on indoor semantic scene segmentation. In order to adapt the network to instance segmentation, we re-implemented the learnable margin method proposed by . The learnable margin method does well on common image instance segmentation datasets and is well-balanced both in speed and accuracy. This combination of two recent papers gives a strong baseline, but not state-of-the-art results in the metrics. We compare to this baseline, also using MinkowskiNet as feature backbone to make the results directly comparable.

We report the average precision (AP) in Table 3 and compare our method with other leading results on ScanNet. Although we do not have the overall state-of-the-art results, the improvement over the baseline verifies the impact of probabilistic embedding and demonstrates that our method can be integrated with any embedding-based method and any backbone network. We can improve the validation mAP by $4.9\%$. We would also like to note that the main point of the paper is to showcase the benefit of the probabilisitc embedding method. We did not have the resources to fine-tune our method for the ScanNet dataset, but nevertheless our results are comparable with the state of the art and in some categories beating state of the art already. We therefore argue that this result underlines the significance of the proposed embedding method as it is likely that future state of the art methods will also be able to benefit from it.

2 Ablation study and analysis

We conduct the ablation study on all categories of PartNet , but we only list detailed values for the four largest categories in Table 4.

We compare four different versions of probabilistic embedding. The Gaussian distribution used in the model can either be isotropic or anisotropic, homoscedastic or heteroscedastic. Thus we have isotropic homoscedastic, anisotropic homoscedastic, isotropic heteroscedastic, and anisotropic heteroscedastic.

The isotropic homoscedastic probabilistic embedding, learns to predict a single scalar representing the uncertainty of a point cloud. We do not see improvements over its determinisitc counterpart, but there is a large gap between them in large categories which have much more part instances and classes than others.

Similar cases happen in anisotropic homoscedastic and isotropic homoscedastic embedding. The former learns a 3D uncertainty vector for a single point cloud, while the latter learns point-dependent uncertainty scalars. They all show significant improvements over determinisitc embedding on fine-grained categories.

Finally, our full model uses anisotropic heteroscedastic probabilistic embedding, which outputs not only point-dependent but also axis-dependent uncertainties. See Figure 7 for an illustration of learned uncertainties. The points at boundary regions have significantly larger uncertainties compared to others. In summary, the full model achieves the best results among all variations.

Effect of spatial embedding.

Since our full model outputs a 3D center vector and 3D uncertainty vector, in a way, we can regard it as a 6D embedding method (with a totally different similarity kernel). One may wonder: how does it compare with the performance of 6D deterministic embedding? The results in Table 4 show, increasing the dimension of deterministic embedding from 3 to 6 shows some improvement, but less than using probabilistic embedding. Thus the performance of our method, cannot be achieved by simply increasing the dimension of deterministic embedding, which also shows the superiority of the probabilistic embedding. We illustrate the differences between deterministic and probabilistic embedding in 3D in Figure 8. We can observe, that probabilistic embedding introduces much stronger deformations of the geometry.

Effect of center-aware loss.

We examine the effect of the center-aware loss in the clustering step. We use the same setup as in our full model except changing the center-aware loss to MSE loss . In Table 4, we can see that our proposed loss function is especially stable on large fine-grained datasets (5.2% vs -7.3%).

Conclusion

We build on embedding-based instance segmentation to present a framework of probabilistic embedding and a new loss function for the clustering step. We evaluate our framework on a large scale point cloud dataset, PartNet, and achieve state-of-the-art performance. Moreover, the qualitative results show the new framework is robust to point clouds with many instances. Additionally, it is able to estimate uncertainties while increasing the accuracy of instance segmentation. In future work, we hope that the probabilistic embedding can be further applied to other kinds of data representation, \eg, 2D images, 3D volumes, and meshes.

Acknowledgements

This work was supported by the KAUST Office of Sponsored Research (OSR) under Award No. OSR-CRG2017-3426.

References

Network architecture

Following the notation of PointNet++ , we give the architecture of the feature network:

where $SA$ and $FP$ are set abstraction and feature propagation module in PointNet++ . The output head network is:

Implementation

We implemented our method using PyTorch and the geometric deep learning library PyTorch Geometric . The final objective function is

Results with different IoU thresholds

We report detailed results of IoU using thresholds of 25% and 75% in Table 5 and Table 6. The metric is mean Average Precision (mAP).

Qualitive Results

We present more qualitive results in Figure 9 which shows the instance-awareness of our method. We also demonstrate the 3D models in the attached video.

Differences to learnable margin

proposed to use a learnable margin for image instance segmentation, which is similar in formulation to our proposed probabilistic embedding. Although we differ in several aspects:

The intuition behind learnable margin comes from the hinge loss: to give different hinge margin to objects of different sizes. However, our intuition comes from modeling neural network outputs as random variables to estimate uncertainty.

The parameters have a different meaning in our method compared to . In learnable margin, $\sigma$ is an instance-specific bandwidth (or margin) per cluster. In our work $\sigma$ are uncertainties per point.

The bandwidth $\sigma$ is influenced by the size of instances (large instances have large $\sigma$). In contrast, our uncertainty $\sigma$ encodes per-point uncertainty close to the boundary of instances (see Fig. 7).

add a loss term to enforce the bandwidths from the same instance to be close. By contrast, we don’t have this kind of restriction. Also, uncertaintes from the same instance can be different as along as they have similar spatial embeddings.