A class-level object pose estimation method based on a learnable prior diffusion model

Abstract

The present invention belongs to the field of three-dimensional perception technology, and discloses a class-level object pose estimation method based on a learnable prior diffusion model. In the training process of the present invention, the parameters of the learnable prior features can be dynamically updated, without the need to collect a specific category of three-dimensional models, which effectively improves the network's learning ability for category information. This method introduces the Transformer diffusion model, and combines it with a position encoding module to enhance the network's ability to understand the three-dimensional information of the object, thereby improving the accuracy of the object's pose estimation. For objects with symmetry, the present invention starts from the probability distribution of the object's pose, cleverly deals with the complexity of multiple possible pose solutions caused by symmetry, effectively avoids the interference of symmetry on pose estimation, and significantly improves the network's estimation performance on such objects. The present invention provides an efficient and accurate pose estimation method, which provides reliable technical support for intelligent robot grasping.

Description

Technical Field

[0001] The present invention belongs to the field of three-dimensional perception technology, and specifically relates to a class-level object pose estimation method based on a learnable prior diffusion model. Background Art

[0002] With the rapid development of artificial intelligence and robotics, intelligent robots are playing an increasingly significant role in industrial production and social life. However, in complex real-world scenarios, accurately perceiving three-dimensional objects in the environment has become a critical issue in the development of robotics. Object pose estimation is a key task in robotic perception systems. It involves determining the position and pose of an object in three-dimensional space. Only by accurately understanding the object's pose can robots achieve precise manipulation. For example, before grasping a glass of water, a robot needs to estimate the orientation of the cup's mouth to avoid spilling the liquid during grasping.

[0003] Although depth cameras provide robots with depth images and RGB images of objects, accurately estimating the pose and position of objects from these images still faces multiple challenges. First, there are many types of objects in real-world scenes, and their sizes and shapes vary significantly. This diversity poses great difficulties for pose estimation. Second, some objects are symmetrical, which leads to multiple possible pose solutions for the same observed image, adding complexity to the pose estimation method. Finally, because depth cameras can only obtain depth images of objects from a certain perspective, they cannot capture depth information of the obscured parts of the object. Therefore, object pose estimation based on incomplete information places higher demands on existing technologies.

[0004] In response to the above problems, several patents have proposed solutions. Patent CN202111542551 reports a six-degree-of-freedom object pose estimation method based on symmetry perception, which can perceive specific objects in the environment and eliminate the errors caused by the symmetry of the objects. Patent CN202411360846 reports a six-degree-of-freedom pose estimation method for objects based on a multi-channel attention mechanism, which enhances the network's understanding of local information through a multi-channel attention mechanism. Although these methods can estimate the pose of an object more accurately, their premise is that the three-dimensional model of the object is known. However, when the objects in the real scene are not fixed, it becomes extremely difficult to collect the three-dimensional models of all objects, which limits the application scenarios of such methods.

[0005] Patent CN202310963602 reports a method for adaptive category-level object pose estimation based on graph convolution dual-stream shape priors. This method achieves category-level object pose estimation by learning a set of 3D model data for objects of a specific category. Although this method solves the problem that existing methods can only process specific objects to a certain extent, it still requires the input of a set of 3D model data of a specific category. This method requires the construction of 3D model data before training, and the network cannot dynamically update data during training, resulting in cumbersome operation steps and limited algorithm learning ability of category information.

[0006] The present invention adopts a Transformer diffusion model regulated by learnable prior features, providing an innovative solution for class-level object pose estimation: by introducing learnable prior features, the present invention does not need to construct a set of three-dimensional model data of specific categories of objects in advance. This prior feature can be dynamically updated during the network training process and continuously learn the object category information in the training set, thereby significantly improving the operating efficiency and learning ability of the method; compared with existing methods based on convolutional neural networks and fully connected networks, this method uses a diffusion model based on Transformer. Compared with convolutional neural networks and fully connected networks, this model exhibits stronger generalization ability and higher accuracy; compared with existing pose estimation methods, this method does not directly predict the object pose, but first predicts the probability distribution of the object pose, and further aggregates the candidate poses, effectively solving the uncertainty problem caused by the symmetry of the object. Summary of the Invention

[0007] In response to the shortcomings of existing pose estimation technology, the present invention proposes a class-level object pose estimation method based on a learnable prior diffusion model, which can accurately estimate the pose and size of multiple classes of objects.

[0008] The technical solution of the present invention:

[0009] A class-level object pose estimation method based on a learnable prior diffusion model. The specific steps are as follows:

[0010] Step 1: Build an object pose estimation dataset

[0011] Collect a sufficient number of similar target objects, capture their depth and RGB images, and annotate their pose and size. Then, segment and extract the target objects, and construct a dataset from the segmented and extracted images to provide data support for parameter training of the subsequent object pose estimation model.

[0012] Step 2: Build an object pose estimation model

[0013] (1) Constructing a pose prediction model: The pose prediction model is responsible for predicting the probability distribution of the target object's pose and outputting a set of candidate object poses. The pose prediction model includes a DINO v2 module, a PointNet++ module, a position encoding module, a priori feature module, a pose encoding module, a temporal encoding module, and a noise prediction module.

[0014] DINO v2 module: The input is a segmented RGB image of the target object with 3 input channels, and the output is the image features of the target object with 384 output channels. DINO v2 is a large image model pre-trained on large-scale datasets. It has powerful object semantic information understanding capabilities and is used to improve the processing of RGB images.

[0015] PointNet++ module: The input is the depth image of the target object and the image features of the target object output by the DINO v2 module. First, the depth image is sampled at the farthest point, 1024 key points are retained, and a key point cloud is obtained. At the same time, the image features of the target object are copied 1024 times as the features of the key points. Subsequently, multiple groups of MLP networks and the K-nearest neighbor algorithm are used to extract the local features of each key point, and the farthest point sampling is used on the sampled point cloud containing local features to halve the number of key points. On this basis, the "local feature extraction" and "farthest point sampling" operations are repeated twice, and finally a point cloud feature containing 128 key points is generated, and the feature dimension of each key point is 256.

[0016] Position Encoding Module: Input is the 3D position coordinates of 128 key points corresponding to the point cloud features generated by the PointNet++ module, and output is the position encoding of the point cloud. For any point P in the point cloud, its 3D coordinates are (x, y, z), and the position encoding of the point is calculated according to the following formula:

[0017]

[0018] Where i represents the channel number of the point cloud feature. The calculated position code is added to the point cloud feature to embed the point cloud's position information into the point cloud feature, improving the model's ability to understand the 3D spatial structure.

[0019] Pose encoding module: extracts random pose features through a two-layer MLP network, concatenates the pose-encoded random pose with the point cloud features embedded with position information, and generates noise markers.

[0020] A priori feature module: This module consists of N 256-channel features, where N represents the number of object categories to be estimated. Depending on the input category, the module outputs the corresponding feature. This feature is synchronized with the parameters of other modules during training.

[0021] Time encoding module: The time features are extracted through a two-layer MLP network, and the time features are added to the category features obtained by the prior feature module to generate the conditional features for noise generation.

[0022] Noise Prediction Module: This module consists of multiple AdaLN modules (Adaptive Layer Normalization modules), multiple attention modules, and a pose decoding module. Guided by conditional features, the AdaLN module normalizes the noise markers input from the pose encoding module and dynamically adjusts the feature distribution. The attention module assigns attention weights to focus on key areas of the object and extracts object pose features from the noise markers. The pose decoding module decodes the object pose from the extracted pose features, ultimately obtaining an estimated candidate pose for the object.

[0023] (2) Constructing a pose scoring model: The structure of the pose scoring model is consistent with that of the pose prediction model. The difference is that the pose scoring model predicts the likelihood probability of the candidate pose of the object predicted by the pose prediction model and outputs the corresponding pose score.

[0024] (3) Constructing a pose aggregation model: The pose aggregation model filters the candidate poses of the objects predicted by the pose prediction model based on the pose scores output by the pose scoring model, and removes abnormal poses with low scores. Subsequently, a clustering algorithm is used to classify the retained candidate poses of the objects. The category with the largest number of candidate poses of the objects is considered to be the correct pose type. Finally, the pose aggregation model calculates the average pose of the correct pose type through mean pooling and uses it as the final object pose output by the model.

[0025] (4) Constructing a size prediction model: The size prediction model predicts the length, width, and height of the target object using its point cloud features. The size prediction model uses the point cloud features output by the PointNet++ module in the pose prediction model as input. Through a multi-layer MLP network, it perceives the scaling properties of the target object and ultimately completes the prediction of the 3D size.

[0026] Step 3: Train the object pose estimation model

[0027] The pose prediction model, pose scoring model, and size prediction model in the object pose estimation model all contain learnable parameters and must be trained sequentially in a predetermined order. The pose prediction model should be trained first, followed by fixing its network parameters, and then the pose scoring model and size prediction model should be trained sequentially.

[0028] (1) Training pose prediction model: The pose prediction model is responsible for predicting the candidate poses of the target object from the object image.

[0029] The specific training process is as follows: First, a set of depth images and RGB images of the target object obtained in step 1 is randomly selected and input into the pose prediction model. During training, random noise is continuously added to the object's true pose to generate noisy samples. Then, the noisy samples are input into the model, and the model learns to restore the denoising process from the noisy samples by learning the mean squared error (MSE) loss function. Specifically, the model uses the Adam optimizer to adjust parameters in each training step to reduce the value of the loss function. The above training process is repeated until the value of the loss function tends to stabilize, indicating that the model training has converged. Finally, the parameters of the pose prediction model are fixed for subsequent model training.

[0030] (2) Training pose scoring model: The pose scoring model is used to score candidate poses and thus eliminate abnormal poses.

[0031] The specific training process is as follows: Since the pose scoring model has the same network structure as the pose prediction model, the parameters of the pose prediction model can be used as the initial parameters. Subsequently, the training process is the same as that of the pose prediction model, but the training goal of the pose scoring model is to learn the likelihood probability of the target object pose rather than directly predict the object pose.

[0032] (3) Training size prediction model: The size prediction model is responsible for predicting the length, width, and height of an object.

[0033] The specific training process is as follows: First, set the pose prediction model to inference mode. Then, randomly extract a set of data from the depth image and RGB image of the target object obtained in step 1 and input it into the pose prediction model. The point cloud features of the target object are extracted from its PointNet++ module. Next, the point cloud features are input into the size prediction model, which outputs the 3D size of the target object. The difference between the output object size and the ground-truth object size obtained in step 1 is calculated and used as the training loss function. The training process is repeated until the loss function stabilizes, indicating that the model training has converged.

[0034] Step 4: Run the process online

[0035] When deployed online, first obtain the depth image and RGB image dataset of the target object according to step 1; then, build the object pose estimation model according to steps 2 and 3, and use the dataset obtained in step 1 to train the pose prediction model, pose scoring model, and size prediction model in the model; after training, fix the model parameters and the model can be used for online object pose estimation.

[0036] Whenever the pose of a new object of the same type needs to be estimated, a depth camera is used to acquire a depth image and RGB image of the target object. The image is then segmented using the Mask-RCNN segmentation algorithm. The segmented image is then fed into a pose prediction model to extract the target object's point cloud features. Using the point cloud features as a condition, the randomly generated pose noise is denoised to generate a set of candidate poses for the target object. These candidate poses are then fed into a pose scoring model to calculate their scores. Next, the pose aggregation model is used to filter and cluster the candidate poses and their corresponding scores, outputting the final pose of the target object. Finally, the target object's point cloud features are fed into a size prediction model to predict the target object's three-dimensional size.

[0037] The present invention has the following advantages: (1) The present invention does not rely on a collection of 3D models of similar objects; only a single training run is required to accurately estimate the pose and size of similar objects. (2) The present invention eliminates the need for cumbersome special processing for symmetrical objects, significantly improving the accuracy of pose estimation for symmetrical objects. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] Figure 1 This is a flow chart of the class-level object pose estimation method based on a learnable prior diffusion model of the present invention.

[0039] Figure 2 It is a structural diagram of the posture prediction model of the present invention. DETAILED DESCRIPTION

[0040] The specific implementation of the present invention is further described below in conjunction with the accompanying drawings and technical solutions.

[0041] After completing the complete training of image data of a specific category of objects, the present invention can estimate the pose and three-dimensional size of all objects of the same category. The specific method is as follows:

[0042] First, construct a target object dataset according to step 1. Collect at least 10 objects of each target category and use a scanner to obtain a 3D model of the target category objects. Subsequently, use a depth camera to record a video sequence of the target objects. The video should contain approximately 100,000 depth images and RGB images. Next, annotate the 2D key points on the target object image and match them with the 3D key points on the target object 3D model to calculate the target object's pose. Finally, segment the target object using the Mask-RCNN segmentation algorithm and extract a single object image. If the target category is more common, you can also directly download an open-source class-level object pose estimation dataset.

[0043] Secondly, build the object pose estimation model according to the description in step 2. Figure 1As shown, the specific process is as follows: (1) Constructing the pose prediction model: According to Figure 2 As shown in the structural diagram, first download the pre-trained dinov2_vits14 model, build the DINOv2 module and load its pre-trained parameters, while fixing the parameters to maintain its feature extraction capability. Subsequently, build the PointNet++ module, the position encoding module and the pose encoding module. Connect the above four modules in sequence according to the instructions in step 2 to generate the noise mark of the diffusion model. In addition, build the prior feature module and the time encoding module, add the output structures of the two modules to generate the conditional features of the diffusion model as the input of the subsequent modules. Finally, build the noise prediction module. According to the instructions in step 2, connect multiple AdaLN modules, multiple attention modules and the pose decoding module, input the noise mark and conditional features of the diffusion model, and finally output the candidate pose of the target object. (2) Build the pose scoring model: Its construction method is the same as the object pose estimation model, only adding the scoring output at the end of the model to evaluate the quality of the candidate pose. (3) Build the pose aggregation model and the size prediction model: According to the instructions in step 2, cascade the relevant models. At this point, the construction of all models is completed.

[0044] Finally, train the object pose estimation model as described in step 3. The trainable components of the object pose estimation model, which contain model parameters, include the pose prediction model, the pose scoring model, and the size prediction model. The pose aggregation model does not require training. The training process must be completed in a strictly sequential manner. The specific process is as follows: First, train the pose prediction model. Then, fix the model parameters as the starting parameters for training the pose scoring model. Finally, set the pose prediction model to inference mode. Input the point cloud features output by the PointNet++ module in the pose prediction model into the size prediction model to train the size prediction model parameters.

[0045] During training, 95,000 images are selected as the training set and 5,000 images as the validation set to evaluate the model training results. When training with an RTX 3090 GPU, the batch size is set to 72. After training the pose prediction model for approximately 72 hours, the model loss typically converges. Next, the batch size is set to 48, and the pose scoring model is trained for approximately 24 hours. Finally, the batch size is set to 128, and the size prediction model is trained for approximately 24 hours. This completes the training of all models. After training, the validation set is used to test the model's performance. If the test results are unsatisfactory, the training time can be extended until the results meet the requirements. Based on experience, a complete training task can typically be completed in approximately 5 days using an RTX 3090 GPU.

[0046] After all models are trained, the pose and size of the target object can be predicted according to the online operation process in step 4.

Claims

1. A class-level object pose estimation method based on a learnable prior diffusion model, characterized in that: The specific steps are as follows: Step 1: Build an object pose estimation dataset Collect a sufficient number of similar target objects, capture their depth images and RGB images, and annotate the pose and size of the target objects; Subsequently, the target object is segmented and extracted, and a dataset is constructed from the segmented and extracted images to provide data support for parameter training of the subsequent object pose estimation model; Step 2: Build an object pose estimation model (1) Constructing a pose prediction model: The pose prediction model is responsible for predicting the probability distribution of the target object's pose and outputting a set of candidate object poses. The pose prediction model includes a DINO v2 module, a PointNet++ module, a position encoding module, a priori feature module, a pose encoding module, a temporal encoding module, and a noise prediction module. (2) Constructing a pose scoring model: The structure of the pose scoring model is consistent with that of the pose prediction model. The difference is that the pose scoring model predicts the likelihood probability of the candidate pose of the object predicted by the pose prediction model and outputs the corresponding pose score; (3) Constructing a pose aggregation model: The pose aggregation model filters the candidate poses of objects predicted by the pose prediction model according to the pose scores output by the pose scoring model, and removes abnormal poses with low scores; Subsequently, a clustering algorithm is used to classify the retained object pose candidates, where the category with the largest number of object pose candidates is considered the correct pose type. Finally, the pose aggregation model calculates the average pose of the correct pose type through mean pooling and uses it as the final object pose output by the model. (4) Constructing a size prediction model: The size prediction model predicts the length, width, and height of the target object through the point cloud features of the target object. The size prediction model uses the point cloud features output by the PointNet++ module in the pose prediction model as input, and through a multi-layer MLP network, it perceives the scaling properties of the target object and finally completes the three-dimensional size prediction. Step 3: Train the object pose estimation model The pose prediction model, pose scoring model, and size prediction model in the object pose estimation model all contain learnable parameters and need to be trained in a predetermined order. The pose prediction model should be trained first, and then its network parameters should be fixed, and then the pose scoring model and size prediction model should be trained in sequence. Step 4: Run the process online When deployed online, first obtain the depth image and RGB image dataset of the target object according to step 1; then, build the object pose estimation model according to steps 2 and 3, and use the dataset obtained in step 1 to train the pose prediction model, pose scoring model, and size prediction model in the model; after training, fix the model parameters and the model can be used for online object pose estimation.

2. The method for class-level object pose estimation based on a learnable prior diffusion model according to claim 1, characterized in that: The pose prediction model is as follows: DINO v2 module: The input is the RGB image of the target object after segmentation, the number of input channels is 3, and the output is the image features of the target object, with 384 output channels; DINO v2 is a large image model pre-trained on a large dataset. It has strong object semantic understanding capabilities and is used to improve the processing of RGB images. PointNet++ module: Input is the depth image of the target object and the image features of the target object output by the DINO v2 module. First, the depth image is sampled at the farthest point, retaining 1024 key points to obtain a key point cloud. At the same time, the image features of the target object are replicated 1024 times as the key point features. Subsequently, multiple MLP networks and the K-nearest neighbor algorithm are used to extract local features for each key point. The sampled point cloud containing local features is sampled at the farthest point to halve the number of key points. On this basis, the "local feature extraction" and "farthest point sampling" operations are repeated twice, ultimately generating a point cloud feature containing 128 key points, each with a feature dimension of 256. Position Encoding Module: Input is the 3D position coordinates of 128 key points corresponding to the point cloud features generated by the PointNet++ module, and output is the position encoding of the point cloud. For any point P in the point cloud, its 3D coordinates are (x, y, z), and the position encoding of the point is calculated according to the following formula: Where i represents the channel number of the point cloud feature; the calculated position code is added to the point cloud feature to embed the point cloud position information into the point cloud feature, improving the model's ability to understand the three-dimensional spatial structure; Pose encoding module: extracts random pose features through a two-layer MLP network, concatenates the pose-encoded random poses with the point cloud features embedded with position information, and generates noise markers; Prior feature module: It consists of N 256-channel features, where N represents the number of categories of objects to be estimated. Depending on the input category, the prior feature module outputs the features of the corresponding category. The parameters of this feature are updated synchronously with other modules during the training process. Time encoding module: extracts time features through a two-layer MLP network, adds the time features to the category features obtained by the prior feature module to generate conditional features for noise generation; Noise prediction module: consists of multiple AdaLN modules, multiple attention modules and a pose decoding module; the AdaLN module normalizes the noise tags input from the pose encoding module under the guidance of conditional features and dynamically adjusts the distribution of features; the attention module focuses on the key areas of the object by assigning attention weights and extracts the object pose features from the noise tags; the pose decoding module decodes the object pose from the extracted object pose features and finally obtains the estimated object candidate pose.

3. The class-level object pose estimation method based on a learnable prior diffusion model according to claim 1, characterized in that: Step 3 is as follows: (1) Training the pose prediction model: The pose prediction model is responsible for predicting the candidate poses of the target object from the object image; The specific training process is as follows: First, a set of depth images and RGB images of the target object obtained in step 1 is randomly selected and input into the pose prediction model. During the training process, random noise is continuously added to the true pose of the object to generate noisy samples. Then, the noisy samples are input into the model, and the model learns to restore the denoising process from the noisy samples by learning the mean squared error (MSE) loss function. Specifically, the model uses the Adam optimizer to adjust the parameters in each training step to reduce the value of the loss function. The above training process is repeated until the value of the loss function tends to be stable, indicating that the model training has converged. Finally, the parameters of the pose prediction model are fixed for subsequent model training. (2) Training the pose scoring model: The pose scoring model is used to score candidate poses and eliminate abnormal poses; The specific training process is as follows: Since the network structure of the pose scoring model is the same as that of the pose prediction model, the parameters of the pose prediction model can be used as the initial parameters. Subsequently, the training process is the same as that of the pose prediction model, but the training goal of the pose scoring model is to learn the likelihood probability of the target object pose rather than directly predict the object pose. (3) Training the size prediction model: The size prediction model is responsible for predicting the length, width, and height of the object; The specific training process is as follows: first, set the pose prediction model to inference mode, then randomly extract a set of data from the depth image and RGB image of the target object obtained in step 1, input it into the pose prediction model, and extract the point cloud features of the target object from its PointNet++ module; then, input the point cloud features into the size prediction model, and output the three-dimensional size of the target object; calculate the difference between the output object size and the true value of the object size obtained in step 1, and use it as the training loss function; repeat the above training process until the value of the loss function tends to stabilize, indicating that the model training has converged.

4. The class-level object pose estimation method based on a learnable prior diffusion model according to claim 1, characterized in that: In step 4: Whenever the pose of a new object of the same type needs to be estimated, the depth image and RGB image of the target object are first acquired through the depth camera, and the image is segmented using the Mask-RCNN segmentation algorithm; then, the segmented image is input into the pose prediction model to extract the point cloud features of the target object; Using point cloud features as conditions, the randomly generated pose noise is denoised to generate a set of candidate poses of the target object; the candidate poses are then input into the pose scoring model to calculate the pose scores of the candidate poses. Scoring; then, combining the candidate poses and corresponding scores, the pose aggregation model is used for screening and clustering. Output the final pose of the target object; finally, input the point cloud features of the target object into the size prediction model, Predict the 3D size of the target object.

Images