Object-Pose Estimation With Neural Population Codes

Pose ambiguity can be dealt with by estimating a probability distribution over the group of rotations in 3D, ${\bf R}\in{\bf SO(3)}$ [7, 8]. The distribution may be represented parametrically, e.g., with a Gaussian mixture model [7] or non-parametrically [8]. The former has the disadvantage of being hard to train [7], e.g., due to local minima [1].

In the non-parametric method by Murphy et al., the distribution is predicted implicitly, i.e., a neural network computes the probability $p({\bf R},I)$ for any input matrix ${\bf R}$ and input image $I$ [8]. The disadvantage of this strategy is that $p$ has to be normalized across a large sample of rotation matrices that all have to be fed through the network, resulting in a computationally expensive and slow training.

Given a normalized probability and a ground-truth rotation matrix ${\bf R}_{o}$, the loss function is defined as follows [8]

$$\mathcal{L}=-\log(p({\bf R}_{o}|I))\,.$$

(1)

For inference, the rotation with maximum probability is chosen from a set of centers of an equivolumetric partitioning of SO(3) [8],

$${\bf\hat{R}}=\underset{{\bf R}\in{\bf SO(3)}}{\arg\max}\,p({\bf R}|I)\,.$$

(2)

Manhardt et al. train a network to predict multiple pose hypotheses to learn alternative (ambiguous) poses [9]. The network output is a fixed number $M$ of possible rotations. For training, the loss function combines a minimum loss $\mathcal{L}_{\mbox{\small min}}$, the minimum mean squared error (MSE) between the ground truth and the $M$ outcomes, with an average loss $\mathcal{L}_{\mbox{\small avg}}$, which is the average MSE across the $M$ outcomes,

$$\mathcal{L}=(1-\epsilon\frac{M}{M-1})\mathcal{L}_{\mbox{\small min}}+\epsilon\frac{M}{M-1}\mathcal{L}_{\mbox{\small avg}}\,,$$

(3)

where the parameter $\epsilon$ is gradually reduced from 0.05 to 0.01.

Inference is more complicated since we do not know a probability for each of the $M$ outcomes. Instead, the outcomes are clustered, and the median for each cluster is computed. The results can be compared against the input image to find the best fit [9]. Though this strategy involves again an extra selection step as mentioned above.

To encode the orientation with a neural population code, we first transform the rotation matrix into a rotation axis and a rotation angle. For the axis, the preferred values of each neuron are arranged on a sphere. To obtain a near uniform distribution of axes, we compute a Fibonacci lattice on a sphere [10]. For the rotation angle, the preferred values are arranged in a circle. To combine axis with angle, our population code has for each point on the Fibonacci sphere a circle of neurons for the rotation angle. So, the space to represent the orientation is a direct product of a sphere and a circle.

Given a rotation axis and angle, we activate all neurons $\{i\}$ using Gaussian tuning curves,

$$a_{i}=\exp\left(-\frac{d\theta_{i}^{2}+d\phi_{i}^{2}}{2\sigma^{2}}\right)\,,$$

(4)

where $d\theta_{i}$ is the angle between the encoded axis and the preferred axis on the sphere for neuron $i$, $d\phi_{i}$ the angle between the encoded angle and the preferred angle of neuron $i$, and $\sigma$ is the tuning width, here $20^{\circ}$.

To compute the target population code from a ground truth transformation matrix, ${\bf R}_{o}$, we compute activations for all symmetry transformations of ${\bf R}_{o}$,

$${\bf R^{\prime}}_{k}={\bf R}_{o}{\bf R}_{\mbox{\small sym}}^{k}\,,$$

(5)

where $\{{\bf R}_{\mbox{\small sym}}^{k}\}$ is the set of symmetry transformations for a given object (in a manufacturing setting, the symmetry of an object is known a priori). The target population code is the sum of all activations $a_{i}$ over all $k$ symmetry transformations (Fig. 1).

The transformation from rotation matrix to axis/angle, ${\bf r}$/$\phi$, is not unique since $-{\bf r}$ and $2\pi-\phi$ is an equivalent axis/angle combination. Therefore, for each symmetry transformation, we compute the tuning-curve activations also for $-{\bf r}$ and $2\pi-\phi$ and add those to the target activations (resulting in four peaks in Fig. 1).

For objects with a symmetry axis, which would result in infinitely many equivalent poses, we compute a variant of our population code: since the population code for the rotation angle would show uniform activations, we simplify the code and omit the rotation angle neurons and just use the neurons encoding the rotation axis on the Fibonacci sphere. Here, the code is computed simply as

$$a_{i}=\exp\left(-\frac{d\theta_{i}^{2}}{2\sigma^{2}}\right)\,,$$

(6)

without having to superimpose activations for symmetries. So, our population code can leverage the rotation symmetry.

For inference, we predict the entire population code and choose the neuron with the highest activation,

$$j=\underset{i}{\arg\max}\,a_{i}\,.$$

(7)

Neuron $j$ corresponds to a rotation axis and angle, which we convert into a rotation matrix. For objects with symmetry axis, we pick a random angle. This inference is similar to (2) from the probabilistic method [8] with the crucial difference that we only need to evaluate our network once, resulting in a faster pose estimate.

Our neural network maps gray-scale images of size 128x128 pixels onto a population code, which is a vector of size $nm$, where $n$ is the number of preferred axes and $m$ the number of preferred angles. Here, we used $n=2,562$ and $m=36$, matching the number of rotation candidates used in [4].

The images are fed through four convolutional layers (see Fig. 2 regarding the kernel size and number of features). The first three convolutional layers are followed by batch normalization and ReLU activation functions. This architectural choice was inspired by the Augmented Autoencoder [4]. The fourth convolutional layer with kernel size 1x1 is followed by a GeLU non-linear function. This architectural element was inspired by [11], which demonstrated its benefit.

After the convolutional layers, the features are mapped through three hidden linear layers onto the population code, which itself acts like a linear layer. Each hidden layer is followed by a Leaky ReLU activation function.

The synthetic-data experiment demonstrates the benefit of the population code. Here, we predict the one-dimensional orientation of shapes inside 64x64 pixels images (Fig. 3). As training data, we used either a set of randomly generated bar or arrow images (800/200 train/test split).

For the synthetic data, we used a simplified neural network comprising three convolutional layers (feature sizes: 32, 64, and 128) and 4 linear layers. All layers except for the output layer were followed by a ReLU activation function, and for the convolutional layers, we added max pooling (2x2).

For the output, we used either a population code vector of 36 neurons or a scalar orientation angle. The population code had preferred angles uniformly distributed in a circle in increments of $10^{\circ}$ and Gaussian tuning curves with width $20^{\circ}$. For the one-hot vector, we use the same output layer as for the population code.

We trained the network with MSE loss for single variable and population code. For the one-hot vector, we tested both MSE and Cross Entropy loss. In addition, tor training the population code, we compared two strategies: 1) having a single activation peak at the ground truth angle and 2) having two activation peaks $180^{\circ}$ apart for shapes with $180^{\circ}$ symmetry.

We repeated train and test runs 10 times, including new train/test splits for each run. As a result, learning the population code with symmetry (two peaks) had the lowest error (Tab. I). Omitting symmetry during training resulted in a slightly larger error, but was not catastrophic. In contrast, learning failed when mapping directly onto the orientation angle for the symmetric bar.

For the real-world data, we computed six different metrics to evaluate the accuracy of estimating pose; specifically, we evaluated the rotation error, assuming a perfect translation estimate. We chose metrics that are suitable for symmetric objects. The first three metrics are used by the BOP: Benchmark for 6D Object Pose Estimation [12]. The remaining three have been also used in the literature.

Given a set of model vertices, $V_{M}$, the MSSD error is computed as

$$e_{\mbox{\small MSSD}}=\underset{{\bf R}_{\mbox{\tiny sym}}\in S_{M}}{\min}\underset{{\bf x}\in V_{M}}{\max}||{\bf\hat{R}}{\bf x}-{\bf R^{\prime}}_{o}{\bf x}||_{2}\,,$$

(8)

where $S_{M}$ is a set of symmetry transforms, s. th., ${\bf R^{\prime}}_{o}={\bf R}_{o}{\bf R}_{\mbox{\small sym}}$. The MSSD accuracy is the mean recall rate of $e_{\mbox{\small MSSD}}<\theta$ with $\theta$ ranging from 5% to 50% of the object diameter in steps of 5% [12].

The MSPD error is similar to the MSSD error, except that it takes into account only the projection onto the image plane,

$$e_{\mbox{\small MSPD}}=\underset{{\bf R}_{\mbox{\tiny sym}}\in S_{M}}{\min}\underset{{\bf x}\in V_{M}}{\max}||{\bf P}({\bf\hat{R}}{\bf x}-{\bf R^{\prime}}_{o}{\bf x})||_{2}\,,$$

(9)

where ${\bf P}$ is the projection onto the image plane. The MSPD accuracy is the average recall rate of $e_{\mbox{\small MSPD}}<\theta$ with $\theta$ ranging from $5w/640$ to $50w/640$ in steps of $5w/640$, where $w$ is the image width in pixels [12].

The VSD error measures the mismatch in appearance between the model in the ground-truth pose and the same model in the estimated pose,

$$e_{\mbox{\small VSD}}=\underset{x\in\hat{V}\cup V_{o}}{\mbox{avg}}\left\{\begin{array}[]{ll}0&\text{if }x\in\hat{V}\cap V_{o}\land|\hat{d}(x)-d_{o}(x)|<\tau\\ 1&\text{otherwise}\,,\end{array}\right.$$

(10)

where $\hat{V}$ and $V_{o}$ are the projected shapes of the model onto the image plane for estimate and ground truth, and $\hat{d}(x)$ and $d_{o}(x)$ are the estimated and ground-truth depths of the model at pixel $x$. The VSD accuracy is the average recall rate of $e_{\mbox{\small VSD}}<\theta$ with parameter $\tau$ ranging from 5% to 50% of the object diameter in steps of 5% and $\theta$ ranging from 0.05 to 0.5 in steps of 0.05 [12].

In this variant, the accuracy is the ratio of errors below threshold 0.3 with $\tau=20$. We adopt this metric since it was used by Sundermeyer et al for the augmented autoencoder [4].

VSS is a variant of VSD where $\tau=\infty$, i.e., the error in depth is ignored. This metric has been used for the multiple pose hypotheses method [9].

ADI is a variant of the Average Distance of Model Points metric, adapted for symmetric objects [13]. The minimization takes into account ambiguous poses.

$$e_{\mbox{\small ADI}}=\underset{{\bf x}\in V_{M}}{\mbox{avg}}\underset{{\bf y}\in V_{M}}{\min}||{\bf\hat{R}}{\bf y}-{\bf R}_{o}{\bf x}||_{2}\,,$$

(11)

To compute the ADI accuracy, we computed the ratio of $e_{\mbox{\small ADI}}<\theta$ with $\theta$ equal to 5% of the object diameter. Here, we used a different threshold than in [13], who used 15%, because with $\theta=15$% of diameter, all our experiments would show a 100% accuracy.

We found large differences in sensitivity to rotation error between metrics (Fig. 4). Particularly, for roughly spherical objects most metrics failed to reflect large errors. Moreover, only MSSD and ADI are invariant to the viewpoint since the other metrics compute a projection onto the image plane. To conclude, we found MSSD to be the most suitable metric to estimate rotation accuracy.

On the real-world data, we compared our method with three other methods. All of them used the same network architecture, as described in Section III, and the same data augmentation. This ensures a direct comparison of the impact of the representations. In contrast, results across publications are difficult to compare because architecture and data augmentation significantly affect the outcomes but are often hard to replicate due to a lack of detailed information. This section describes the alternative methods.

We adapted the multiple pose hypotheses method [9] to our network, predicting $M$ quaternions in the last layer (size: $4M$). For training, we used the loss in (3), and for inference, we computed mean-shift clustering on the $M$ predicted quaternions [9] and picked the mean of the largest cluster.

For objects with a symmetry axis, we can directly predict the three coordinates of the axis since there is no ambiguity. Here, the last layer of our network had only three elements, and for training, we used the MSE to the ground truth axis.

For objects with discrete symmetries, we computed the minimum loss to the closest target (one for each equivalent pose). Here, the last layer of our network was 6 dimensional for storing two columns of the rotation matrix R. Learning two columns has been shown to be advantageous over other representations like quaternions due to the continuity of the R6 space [14]. This representation was also used by one of the best performing methods for asymmetric objects [15], and, as we will see, for asymmetric objects, mapping onto R6 does indeed work well (Tab. II). The entire rotation matrix can be reconstructed from the two columns using Gram-Schmidt orthonormalization. For training, we used the L1 loss.

Instead of treating the output layer as a population code, we used it as a one-hot vector, where each element/index represents a rotation. As common for one-hot vectors, we used the Cross-Entropy loss for training. For inference, we picked the index with the largest logit/probability.

We carried out our real-world experiments with the T-LESS dataset [16], which contains feature-less and mostly symmetric objects. Table II shows all 30 objects.

For training, we used the images from the BlenderProc4BOP dataset [17], extracting all instances of an object with at least 60% visibility. We augmented these data with the T-LESS Primesense camera images [16]. For all images, we used the provided bounding boxes and made square cutouts by using the maximum of width and height plus a 10% padding. The resulting cutouts were scaled to 128x128 pixels. We supplemented these images with 8,000 generated images per object using pyrender and the 3D model files from the T-LESS dataset. This combination resulted in about 20,000 to 22,000 images per object.

We combined some pairs of visually similar objects into one dataset and trained a single model for both objects (Tab. II). This strategy resulted a small improvement (Tab. VI) and also points to the ability of our architecture and population-code representation to generalize across objects.

We trained all methods for 80 epochs, using the Adam optimizer with a learning rate of 0.0002 and a batch size of 64. As the loss function, we used MSE, except as specified differently above (Section IV.C). During training, we augmented the data for the pyrender-generated images by 1) adding random brightness, a uniformly-distributed value between -0.2 and 0.2 to all pixels with a 50% probability and 2) adding a background image with a 70% probability. Moreover, for all training images, we added Gaussian noise (std: 0.02) with a 50% probability, did contrast normalization, and shifted the image in the image plane ($\pm 5$ pixels in x and y) with 90% probability.

For testing, we used the official test set from the BOP Pose Estimation Challenge [17]. Since our focus was on rotation estimation rather than object detection, we used the provided ground truth bounding boxes. We compared Single Variable, One-Hot Vector, and Population Code across all 30 objects (Tab. II), whereas Multiple Pose Hypotheses was evaluated only on objects 4 and 5, as its performance was no better than random weights (Tab. IV and V). This suggests that the approach cannot learn from scratch and requires a pre-trained network backbone, as used by [9].

For our method, the average inference time was 3.2 msec on a Macbook Air M1 2020, which makes our method suitable for edge devices and real-time applications. In comparison, Multiple Pose Hypotheses required 115 msec for inference due to the costly clustering.

For the population code, the average MSSD accuracy across all objects was 84.7% (Tab. III). This value compares favorably against the best result on the BOP leaderboard for RGB-only input in the category Localization of Seen Objects: 81.4% (https://bop.felk.cvut.cz/leaderboards/pose-estimation-bop19/t-less/ accessed on 9/14/2024). The limitation is that we did not estimate the translation but instead used the ground truth translation.

Our average $e_{\mbox{\tiny VSD}}<0.3$ was 87.3% across all objects (Tab. III). In comparison, using the same ground-truth bounding boxes, Sundermeyer et al reported a 38.34% accuracy from RGB input and 84.04% with ICP refinement on the depth data [4]. However, the same caveat as above holds regarding estimating the translation.

Predicting a single variable performed well when there was no ambiguity in the target (Tab. III) but struggled with ambiguous poses (Tab. III). In contrast, our population code method performed equally well under ambiguity (Tab. III). Moreover, population code targets consistently outperformed one-hot vectors by large margins (Tab. II).

We also observed that our network faithfully learned the entire distribution of the population code (Fig. 5) and not just the peaks. This behavior is expected since we compute the MSE across all neurons of a population code. Interestingly, even without providing symmetry information during training, the performance drop was relatively small (Tab. VI); i.e., the network could still learn to represent the ambiguity caused by symmetry.