Single-Stage 6D Object Pose Estimation

In plain words

To find where an object is and how it is turned, many systems first find image spots that match fixed points on the object's 3D model, then run a classic geometric solver with random sampling to fit the pose to those spots. The sampling step is slow, and the network that finds the spots is trained on a proxy error, so the final pose can be off even when the spots look fine.

This work replaces the geometric solver with a small neural network. For each model point, the network receives a whole bunch of candidate spots. Inside a bunch the order of the candidates does not matter, so the network pools them. The order of the model points does matter, so each bunch keeps its own slot in the input. The network then outputs the pose directly and is trained on a pose error. Because it is only a replacement for the solver, it can be attached to existing spot-finding networks and trained with them.

TL;DR

• Replaces RANSAC-based PnP in two-stage 6D pose pipelines with a learned, order-aware network §1 §3.1.
• Input: for each of n fixed 3D keypoints, m candidate 2D correspondences; output: a quaternion and a translation §3.1.
• Used with SegDriven and PVNet correspondence networks, it raises the Occluded-LINEMOD ADD-0.1d average and is faster than PVNet's pipeline on the same backbone tab. 1 tab. 2 tab. 3.
• Results are author-reported; the released repository contains only the pose network, not the real-data pipeline README.md.

Problem

Two-stage pipelines learn 2D correspondences with a surrogate loss on 2D error, then solve PnP with RANSAC §1. The paper argues that equal average 2D error can produce different poses, and that RANSAC output depends on the order of correspondences fig. 1. The pipeline is not end-to-end trainable, and RANSAC is slow when there are many outliers §1.

Method

Problem formulation §3.1. Given the camera matrix K and, for each of n 3D keypoints p_i, m candidate 2D points u_ik, the pose (R, t) satisfies the projection relation eq. 1. The method learns a function g with parameters Θ that maps all (p_i, u_ik) pairs directly to (R, t) eq. 2. Classical PnP solvers instead use RANSAC over random subsets §3.1.

Properties the network must respect §3.1.1.

• Cluster ordering: the order of candidates inside one keypoint's cluster must not change the output, but the order of the clusters is fixed by the keypoints §3.1.1.
• Interaction: one cluster alone says nothing about the pose; the global structure across clusters must be captured §3.1.1.
• Absolute positions: unlike point-cloud networks, the features must depend on absolute image positions, not on rigid transformations §3.1.1.

Architecture §3.1.2 fig. 3.

1. Local features: a three-layer MLP with weights shared across all candidates and clusters produces a D-dimensional feature f_ik for each candidate §3.1.2.
2. Grouped aggregation: for each keypoint i, take the element-wise max over its m candidate features; concatenate the n results in keypoint order to get an nD vector eq. 3 §3.1.2.
3. No instance or batch normalisation is used; the paper reports that neither helped §3.1.2.
4. Global inference: three fully connected layers map the nD vector to a quaternion and a translation §3.1.2.

The paper contrasts this with a single global max-pool over all points, which would be permutation invariant but would ignore the fixed keypoint order §3.1.2.

Single-stage framework §3.2. A second network f(p_i, I; Φ) outputs the m candidate correspondences for keypoint i from the image I eq. 4. Concretely, f is the encoder-decoder of SegDriven [Hu19a] or PVNet [Peng19a] §3.2. The complete model composes f over all keypoints with g eq. 5.

• A single cube is used for all objects, so its eight corners are the keypoints. The cube is the largest one inside a sphere whose radius is the average of the object bounding spheres §3.2. This means the keypoint identity is implied by the cluster order, not supplied as input §3.2.
• Input per candidate is 4D: the grid-cell centre (x, y) and an offset (dx, dy). The image point is x + dx, y + dy. The paper reports that feeding the image point directly worked worse than this first-order form §3.2. With PVNet's direction-vector output, dx and dy are normalised to unit direction §3.2.
• m = 200 grid cells are sampled inside the predicted mask of the class §3.2.

Training loss §3.2. The total loss is L = L_s + L_k + L_p eq. 6. L_s is the focal loss for per-cell class labels [Lin17]. L_k is the keypoint regression term of the chosen backbone. L_p is the mean, over the n keypoints, of the distance between the 3D keypoints under the predicted and the ground-truth pose eq. 7. Rotations come from quaternions in a differentiable way §3.2. Translations are normalised so their regression targets have a comparable range §3.2.

# Pose from correspondence clusters  [§3.1.2] [eq. 3]
# Input:  correspondences (B, n, m, 4): per keypoint, m cells of (x, y, dx, dy)
# Output: quaternion (B, 4), translation (B, 3)

def pose_from_clusters(correspondences):
    # Per-candidate features, one shared MLP for all keypoints  [§3.1.2]
    features = shared_mlp(correspondences)              # (B, n, m, D)
    # Element-wise max over the m candidates of each keypoint  [eq. 3]
    cluster_features = max_over_candidates(features)    # (B, n, D)
    # Keep keypoint order: concatenate in the fixed keypoint sequence  [eq. 3]
    global_vector = flatten(cluster_features)           # (B, n * D)
    # Three fully connected layers give the raw outputs  [§3.1.2]
    raw = global_head(global_vector)                    # (B, 7)
    # First four values form the quaternion  [§3.1.2] [code: train.py#L139-L140]
    quaternion = raw[:, :4]
    # Remaining three form a normalised translation  [§3.2] [code: train.py#L143-L145]
    translation = (raw[:, 4:] + 0.5) * (t_max - t_min) + t_min
    return quaternion, translation

# Pose loss  [eq. 7]
# Input:  keypoints (n, 3); predicted and ground-truth quaternion and translation
# Output: scalar

def pose_loss(keypoints, pred_q, pred_t, gt_q, gt_t):
    # Normalise quaternions before converting them to rotation matrices
    # [code: train.py#L22]
    pred_r = rotation_from_quaternion(pred_q / norm(pred_q))
    gt_r = rotation_from_quaternion(gt_q / norm(gt_q))
    # Move the same keypoints under both poses  [eq. 7]
    pred_points = pred_r @ keypoints.T + pred_t    # (3, n)
    gt_points = gt_r @ keypoints.T + gt_t          # (3, n)
    # Average Euclidean distance over keypoints  [eq. 7]
    return mean(norm(pred_points - gt_points, axis=0))

Hyperparameters

Values in the paper and the released code often differ. Both are listed; the Paper vs code section explains the gaps.

Results

Author-reported. Setup: synthetic data with a unit sphere §4.1; real data on Occluded-LINEMOD (8 objects, 1,214 test images) and YCB-Video (21 objects, about 130K frames) §4.2; baselines PoseCNN, SegDriven and PVNet §4.2. Hardware for timing is a GTX 1080 Ti tab. 3. Table and figure numbers follow the order of the LaTeX source; the compiled PDF was not checked.

Synthetic data §4.1 fig. 6 fig. 7 fig. 8. The paper reports results only as plots. In words: RANSAC-EPnP and RANSAC-P3P perform similarly and beat the learned method at very low noise. The learned method becomes clearly more accurate as noise grows fig. 6. Against PVNet's voting-based PnP, the voting method wins at low noise and the learned method is more robust at high noise fig. 8. A single max-pool without cluster order degrades accuracy markedly fig. 7.

Occluded-LINEMOD, ADD-0.1d and REP-5px tab. 2 (%, percentage of correct poses).

* symmetric objects, scored with the symmetric metric §4.2.

Replacing the RANSAC step on Occluded-LINEMOD, ADD-0.1d tab. 1.

Speed on a GTX 1080 Ti, milliseconds tab. 3. The paper says the PVNet backbone with this network runs about twice as fast as PVNet's own pipeline.

YCB-Video tab. 4 (%).

The "–" marks a value missing from the original PVNet paper tab. 4.

Paper vs code

Official code at commit 8e524666e5d8d19f591461e907b2908048acff0a (GPL-3.0, cited only). It contains the pose network, the synthetic data generator, and training and test scripts. It does not contain the correspondence networks, the real-data pipeline or the Occluded-LINEMOD or YCB experiments README.md.

• Training schedule. Paper: 300 epochs, Adam at lr 1e-3, batch 32 §4.1. Code: 200 epochs, Adam at lr 1e-4, batch 32, with the rate cut by 10 at epochs 100, 160 and 180 train.py.
• Synthetic grid sampling. Paper: 200 cells sampled within the mask §4.1. Code: 200 cells drawn uniformly from the whole image, with no mask dataset.py.
• Noise. Paper: "2D noise with variance σ" in [0, 15] §4.1. Code: a Gaussian with standard deviation σ, σ drawn uniformly in [0, 15] per sample dataset.py.
• Outliers. Paper: a percentage of dx, dy is replaced by uniform values §4.1. Code: the grid position x, y is replaced as well, and the fraction is drawn per sample in [0, 0.3] dataset.py.
• Test protocol. Paper: sweeps noise and outlier rate §4.1. Code: noise 0 to 30 with a fixed 10% outlier fraction, 2,000 samples per level test.py. Only the EPnP baseline is active; the P3P line is commented out test.py.
• Error metric. Paper: error as a ratio to the target object's diameter §4.1. Code: the mean keypoint distance in model units, with no division by diameter test.py.
• Architecture. Paper: a three-layer local MLP §3.1.2. Code: the third layer is linear, with no activation, before pooling model.py. The global head uses LeakyReLU between its hidden layers and a linear output model.py.
• Keypoint count. Code hard-codes eight keypoints in the reshape model.py; the paper uses eight corners §3.2.
• Translation decoding. Not stated in the paper beyond "normalised" §3.2. Code maps the raw output to the synthetic range [-2, 2] × [-2, 2] × [4, 8] train.py.
• Loss. Paper: L = L_s + L_k + L_p eq. 6. Code: only L_p is computed, with alpha = 1 train.py.
• Consistent points. Eight keypoints, the three-layer global head, and no normalisation layers all match the paper model.py.

No discrepancies were found for the reported real-data numbers, because the released code does not run those experiments.

Pitfalls

• Data layout is silent. The code pools by reshaping the flat candidate list into (batch, channels, candidates, 8). This requires keypoint index as the innermost axis and candidates grouped by grid cell. A different layout still runs and returns a number, but it pools across the wrong axis dataset.py model.py.
• Keypoint count and order are fixed. The paper says the network works only for a fixed set of 3D points §4.3. Changing keypoints requires retraining. Eight is hard-coded in the released model model.py.
• Quaternion scale. The network output is not normalised; the loss normalises the quaternion first train.py. Gradients pass through that normalisation. The quaternion sign is not fixed by the loss, since q and −q give the same rotation matrix (an inference from the code, not stated in the paper).
• Translation range. Predicted translations are confined to a box chosen from the synthetic setup train.py. Real-data ranges are not stated in the paper.
• Input normalisation is not in the paper. The released code centres coordinates by image size and scales offsets train.py. Without it, the scale of the inputs differs.
• Variance versus standard deviation. The paper's "variance σ" is a standard deviation in the code dataset.py. This changes the noise levels by a square.
• Epoch semantics. Each epoch draws fresh random poses and noise, so the "20K images" are resampled every epoch dataset.py train.py.
• Symmetric objects. Symmetric objects are scored with a symmetric metric, and training pose ranges are restricted to avoid symmetry confusion §4.2. The released code has no symmetry handling.
• Speed claims depend on the backbone. The timing comparison uses the PVNet backbone for the extraction step, which takes 14 ms in both rows. The gap comes from fusion, 8 ms against 26 ms tab. 3. The timing hardware is a single GPU.
• Limitation. The network is less accurate than geometric PnP when very precise correspondences are available §4.3 fig. 6.
• Not reproducible from the release. The real-data results need the correspondence networks, which are not in this repository README.md.

Lineage

• Builds on: SegDriven [Hu19a] and PVNet [Peng19a] as correspondence extractors, with the same encoder-decoder §3.2. Classical PnP solvers EPnP [Lepetit09] and P3P [Gao03] are the baselines on synthetic data §4.1. The focal loss [Lin17] is used for segmentation §3.2.
• Compared against: PoseCNN [Xiang18b], SegDriven [Hu19a], PVNet [Peng19a] §4.2.
• Superseded by: not known from content already in this wiki.