Towards Real-World Aerial Vision Guidance with Categorical 6D Pose Tracker

In this subsection, we will present an overview of our proposed category-level 6-DoF pose tracker and then provide detailed introductions to each component in our designed network. As depicted in Fig. 3, our goal is to estimate continuously the change of the 6-DoF pose, denoted as $\Delta{{\cal P}^{(t)}}$, for the target object within an arbitrary known category. The core inputs of our network include the observed RGB-D video stream captured by the onboard camera and the corresponding categorical shape prior ${P_{r}}\in{\mathbb{R}^{{N_{r}}\times 3}}$, which is converted in advance into the same coordinate as the camera. For simplicity, the number of shape prior models is uniformly sampled to be consistent with ${P^{(t)}}$, w.r.t., ${N_{r}}=N$. Different from the recent methods [5, 6, 7, 8, 9] for category-level pose tracking, we employ a three-stage pipline, as displayed in Fig. 3. stage-1: We first integrate the local pixel-point dense feature descriptor for each target object using the proposed 2D-3D Dense Fusion Transformer (Sec. 4.1.2); stage-2: Subsequently, we introduce a Shape-Based Spatial-Temporal Augmentation module with a encoder-decoder structure to dynamically enhance this object-aware descriptor utilizing both temporal prior and shape prior knowledges. It ensures the adaptability of final augmented representations to inter-class variations and inter-frame differences (Sec. 4.1.3); stage-3: All enhanced embeddings are passed through the proposed Prior-Guided Keypoints Generation and Match module to build the 3D-3D inter-frame keypoint pair correspondences in a coarse-to-fine manner (Sec. 4.1.4). The final pose tracking is solved with the Perspective-n-Point (PnP) algorithm and RANSAC using these generated and aligned sets of keypoint pairs.

The objective of this module is to build a local aggregated descriptor for each object by establishing dense per-point feature correspondences between the 3D point patch and the 2D image crop, that serves as the base embeddings for next embedding argumantation. In earlier works such as [53] and [5], 2D image and 3D depth information were used separately as inputs without considering the combination of modal-wise features, that resulted in the loss of intermodal correlation during the feature extraction process. In this regard, we present a pixel-point dense fusion module that utilizes the similarity properties of Transformer to enhance the selection of highly correlated feature pairs, as shown in Fig. 4 (a). Concretely, given the current image pixel crop ${I^{(t)}}\in{\mathbb{R}^{H\times W\times 3}}$, along with the observable geometric point patch ${P^{(t)}}\in{\mathbb{R}^{N\times 3}}$ with a one-to-one correspondence through back-projection, we first employ our proposed Weight-Shared Attention (WSA) to map each pixel in the image crop to a color feature embedding ${F_{c}}\in{\mathbb{R}^{N\times{d_{rgb}}}}$, meanwhile, process the corresponding point in the 3D point patch to a geometric feature embedding ${F_{g}}\in{\mathbb{R}^{N\times{d_{geo}}}}$. The WSA layer adopts an offset-attention structure:

where $\varphi$ represents the linear and ReLU layer applied to the output features and $\alpha$ denotes the softmax function. ${\mathcal{F}_{i}},i=q,k,v$ represents the convolutional operation for query, key and value, respectively.

After the initial dense fusion, we aggregate these base dense information and then encode the context-dependent local feature descriptor $\tilde{F}_{obj}^{(t)}$ for each object in current frame. Inspired by the standpoint proposed by Wang et al. [54], that the low-rank nature of the context mapping matrix in the self-attention mechanism, we utilize this property not only to reduce the complexity time from $O({N^{2}})$ to $O(N)$ but also to enhance the instance’s pose representation in term of local per-point fusion. Specifically, ${F_{c}}$ and ${F_{g}}$ undergo an MLP operation and the color embedding ${F_{c}}$ is projected into two identical projection matrices $X_{c},Y_{c}\in{\mathbb{R}^{N\times k}}$. As shown in Fig. 4 (a), we then incorporate them when computing the key and value vectors. This allows us to calculate an ${(N\times k)}$-dimensional context mapping matrix using scaled dot-product attention with multi-heads:

where ${d}$ is the embedding dimension and $h$ is the number of heads. $\mathcal{F}$ denotes the the linear layer.

Unlike common indoor tabletop scene, the fast change of the onboard camera’s view in the aerial brid’s-eye perspective, such as pitch or roll, may induce the motion blur or significant inter-frame variations in space scale, etc. These challenges are unavoidable in real-time aerial 6-DoF pose tracking. Additionally, the intra-class shape variation with the same class can notably impact the performance of pose tracking for different instances. To our best knowledge, existing category-level pose tracking methods [5, 6, 7, 8, 9] have not completely solved these problems. In this end, we introduce a shape-based spatial-temporal augmentation strategy in this module. This strategy has an encoder-decoder structure leveraging both temporal knowledge dynamic filtering and the shape-similarity filtering, as depicted in the middle of Fig. 3.

The spatial-temporal filtering encoder aims to solve the challenge of inter-frame differences and initially construct a base temporal embedding $\tilde{F}_{temp}^{(t)}$ by transforming the prior knowledge from previous frame to current frame. As presented in the Fig. 4 (b), given the local descriptor $\tilde{F}_{obj}^{(t)}=\{f_{i}^{(t)}\in{\mathbb{R}^{d}}\}_{i={\rm{1}}}^{N}$, we first apply a multi-head attention layer to generate ${{\dot{F}}^{(t)}}=\{\dot{f}_{i}^{(t)}\in{\mathbb{R}^{d}}\}_{i={\rm{1}}}^{N}$:

where the ”${Norm}$” indicates the normalization layer. Considering that a faster change in viewpoint results in fewer overlaps among continous frames, making it difficult to capture useful inter-frame information, we need to retain the observable shape features while filtering out irrelevant data. Therefore, based on $\tilde{F}_{obj}^{(t-{\rm{1}})}=\{f_{j}^{(t-{\rm{1}})}\in{\mathbb{R}^{d}}\}_{j={\rm{1}}}^{N}$ in the previous ${t-1}$-th frame, we compute a spatial-temporal similarity map matrix ${S_{\Delta t}}$ using the vector inner product:

and then is row-aware normalized through a softmax function to constrained the column elements of ${S_{\Delta t}}$ into the range of $[0,{\rm{1}}]$, i.e., ${{\bar{S}}_{\Delta t}}=Softmax{({S_{\Delta t}}(i,\cdot))_{j}}$. Where $<,>$ is the inner product. Afterward, we employ a max-pooling operator along with a convolution layer to obtain the filtered descriptor denoted by ${{\ddot{F}}^{(t)}}=\{\ddot{f}_{i}^{(t)}\in{\mathbb{R}^{d}}\}_{i={\rm{1}}}^{N}$:

where $\mathcal{F}$ denotes the convolution layer and $[;]$ indicates vector concatenation. In this process, we effectively assign weights according to the impact of the previous frame features using the map matrix ${{\bar{S}}_{\Delta t}}$ and prioritize its the most relevant shape points. With this, the current $t$-th temporal embedding can be obtained as follows:

To address another challenge of intra-category shape variability, we then employ the canonical shape prior information to augment obtained temporal embedding in the following augmentation decoder. In this end, a shape-similarity filtering block, as depicted in the bottom of Fig. 4 (c), is adopted before the decoding process to adaptively enhance the local object descriptor based on the shape prior ${P_{r}}$. To jointly optimize the static prior model with the primary network, ${P_{r}}$ is first fed into a learnable layer to generate the shape-point representation ${F_{pr}}=\{f_{j}^{pr}\in{\mathbb{R}^{d}}\}_{j={\rm{1}}}^{{N_{r}}}$. Likewise, a shape-similarity map matrix ${S_{pr}}$ is computed as follows:

and then normalized as ${{\bar{S}}_{pr}}$ using the same operation as before. With this, the filtered features can be denoted by ${{\dddot{F}}^{(t)}}=\{\dddot{f}_{i}^{(t)}\in{\mathbb{R}^{d}}\}_{i={\text{1}}}^{N}$:

where ${\rho_{j}}$ is the coordinate of the shape-point. We assign weights according to the impact of shape-point features using the map matrix ${{\bar{S}}_{pr}}$ to compensate the missing information in current observation. The final augmented output, $\tilde{F}_{aug}^{(t)}$, is updated by adopting one multi-head attention layer with feed-forward, expressed as:

According to these group of augmented representations $\{\tilde{F}_{obj}^{(t)},\tilde{F}_{temp}^{(t)},\tilde{F}_{aug}^{(t)}\}$, we now employ these representations to generate the 3D keypoints for final 6-DoF pose tracking. Different from the prior work 6-PACK [5], where unsupervised keypoint generation may result in a local optimum, we dynamically adapt keypoint generation based on the structural similarity between categorical prior ${P_{r}}$ and observable point patch ${P^{(t)}}$ in the current frame. In the end, we introduce an auxiliary module to convert ${P^{(t)}}$ into $n$ object key-points $[{k_{1}},\ldots,{k_{n}}]$, As shown in the right of Fig. 3, we apply a low-rank Transformer network with $\{\tilde{F}_{obj}^{(t)},\tilde{F}_{temp}^{(t)},\tilde{F}_{aug}^{(t)}\}$ as query, key and value to estimate a structure regularized projection matrix $M\in{\mathbb{R}^{n\times N}}$ for each category. To encourage the key-point transformation $M$ to adapt the intra-class structural variation among different instances, we utilize the shape prior ${P_{r}}$ to optimize this auxiliary module by minimizing the loss ${L_{aux}}$ during training step:

where $P_{r}^{K}=M\times{P_{r}}$ is the prior-base n object keypoints. This formulation effectively ensures that the 3D space of the key-points is structurally consistent with the shape prior, regardless of the pose change over time.

We then apply a 3D tri-plane as a compact feature representation of the projected keypoints, based on the 2D-base formulation in [55]. We align these embedding group $\{\tilde{F}_{obj}^{(t)},\tilde{F}_{temp}^{(t)},\tilde{F}_{aug}^{(t)}\}$ along three axis-aligned orthogonal feature planes by projecting them onto the triplane $\{{T_{XY}},{T_{YZ}},{T_{XZ}}\}$ using the known camera intrinsics. In our implementation, each plane has dimensions $N\times{d_{T}}$. For any object key-point, we project it onto each planes, query the corresponding point feature $\{{T_{xy}},{T_{yx}},{T_{xz}}\}$ via nearest-neighbor point interpolation, which is then conrelated into final keypoint feature $\tilde{F}_{kp}^{(t)}$.

Due to the identity of the projection matrix $M$, a rough match has been revealed between the keypoint pairs between consecutive frames. Futhermore, we then perfrom the finer keypoint matching to filter possible outlier coarse matches. Following [56], a score matrix $S_{kp}$ is calculated based on the simularity between two sets of keypoint features $\tilde{F}_{kp}^{(t-1)}$ and $\tilde{F}_{kp}^{(t)}$ in previous and current frames:

where $\tau$ is a scale factor. We also apply a dual-softmax operator [57] on both dimensions of $S_{kp}$ to obtain the keypoint pairs matching probability:

With the confident matrix ${\mathcal{P}_{c}}$, we select finer key-points with confidence higher than a threshold of ${\theta_{c}}$, and further enforce the mutual nearest neighbor (MNN) criteria:

To improve the performance of our keypoint generation module, we use the following multi-view consistency loss to render the generated keypoints in each of two consecutive frames a better match, placing the keypoint in current view at the transformed corresponding keypoint using ground-truth pose change in previous frame:

Meanwhile, to supervise the matching probability matrix ${\mathcal{P}_{c}}$, we follow LoFTR [56] to use the negative log-likelihood loss over the grids in $\mathcal{M}_{c}^{gt}$. We likewise use camera poses and depth maps to compute the ground-truth for ${\mathcal{P}_{c}}^{gt}$ and $\mathcal{M}_{c}^{gt}$:

The above loss functions only guarantees that the generated keypoint pairs are robust to the change in pose. However, it does not ensure these keypoints are optimal for estimating the final pose. In this regard, we use a differentiable pose tracking loss function, which includes a translation loss and a rotation loss:

Therefore, the overall loss function can be determined as the weighted sum of all losses.

Given the current estimated 3D rotation ${R^{(t)}}$ and the corresponding desired value ${R^{*}}$, we can obtain the change of rotation, i.e., $\Delta R={({R^{(t)}})^{T}}\cdot{R^{*}}$. Let’s use the vector $u\theta$ to express $\Delta R$, where the $u$ represents the rotation axis, and $\theta$ is the rotation angle obtained from identity matrix $\Delta R$. So the objective function for rotational action can be defined as the error of the $\theta{u^{T}}$ toward zero, i.e., ${\varepsilon_{r}}=\theta{u^{T}}-0$, and its time derivative can be related to the camera velocity component generated from onboard manipulator ${V_{C}}^{(m)}$:

where the Jacobian matrix ${L}(u,\theta)$ is

$\sin c(\theta)={{\sin(\theta)}\mathord{\left/{\vphantom{{\sin(\theta)}\theta}}\right.\kern-1.2pt}\theta}$ and ${L(u)}$ is antisymmetric matrix associated with $u$. The error ${\varepsilon_{r}}$ can be converged exponentially by imposing ${\dot{\varepsilon}_{r}}={-\lambda_{r}}{\varepsilon_{r}}$ and $\lambda_{r}$ tunes the convergence rate. Meanwhile, based on the special form of ${L}(u,\theta)$, we can set ${L}(u,\theta)={L}(u,\theta)^{-1}=I_{3\times 3}$ for the small value of $\theta$. Then we compute the relationship between the vector ${V_{C}}^{(m)}$ and the angular velocity vector of each joint of the manipulator $\dot{\eta}$, which can be expressed as: