Key-Grid: Unsupervised 3D Keypoints Detection using Grid Heatmap Features

Chengkai Hou,Zhengrong Xue,Bingyang Zhou,Jinghan Ke,Lin Shao,Huazhe Xu

2024-10-16

Abstract:Detecting 3D keypoints with semantic consistency is widely used in many scenarios such as pose estimation, shape registration and robotics. Currently, most unsupervised 3D keypoint detection methods focus on the rigid-body objects. However, when faced with deformable objects, the keypoints they identify do not preserve semantic consistency well. In this paper, we introduce an innovative unsupervised keypoint detector Key-Grid for both the rigid-body and deformable objects, which is an autoencoder framework. The encoder predicts keypoints and the decoder utilizes the generated keypoints to reconstruct the objects. Unlike previous work, we leverage the identified keypoint in formation to form a 3D grid feature heatmap called grid heatmap, which is used in the decoder section. Grid heatmap is a novel concept that represents the latent variables for grid points sampled uniformly in the 3D cubic space, where these variables are the shortest distance between the grid points and the skeleton connected by keypoint pairs. Meanwhile, we incorporate the information from each layer of the encoder into the decoder section. We conduct an extensive evaluation of Key-Grid on a list of benchmark datasets. Key-Grid achieves the state-of-the-art performance on the semantic consistency and position accuracy of keypoints. Moreover, we demonstrate the robustness of Key-Grid to noise and downsampling. In addition, we achieve SE-(3) invariance of keypoints though generalizing Key-Grid to a SE(3)-invariant backbone.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper attempts to address the problem of detecting 3D keypoints under unsupervised conditions, particularly for deformable objects. Most current unsupervised 3D keypoint detection methods mainly focus on rigid objects, but when dealing with deformable objects, the keypoints detected by these methods often fail to maintain semantic consistency. Therefore, this paper proposes a new unsupervised keypoint detector, Key-Grid, which aims to handle both rigid and deformable objects while maintaining the semantic consistency of keypoints during shape changes.

Specifically, Key-Grid addresses this problem in the following ways:

1. **Autoencoder Framework**: Key-Grid adopts an autoencoder architecture, where the encoder predicts keypoints, and the decoder reconstructs the object using the generated keypoints.
2. **Grid Heatmap**: This is a novel concept used to represent the shortest distance between uniformly sampled grid points in 3D cubic space and the "skeleton" formed by the connections of keypoints. The grid heatmap is used in the decoder part to provide rich geometric descriptions, especially when the object undergoes significant deformation.
3. **Multi-layer Feature Fusion**: Each layer of features in the encoder is directly copied and fused into the corresponding layer of the decoder to further assist the reconstruction process.
4. **Training Objective**: Key-Grid is optimized by minimizing the difference between the reconstructed point cloud and the input point cloud, including similarity loss and farthest point keypoint loss.

Experimental results show that Key-Grid not only achieves state-of-the-art performance on rigid objects but also significantly outperforms existing methods in the detection of deformable objects. Additionally, Key-Grid demonstrates robustness to noise and downsampling and can handle objects in different poses by extending to SE(3)-equivariant versions.