Class level object pose estimation method based on learnable prior diffusion model

Abstract

The invention belongs to the technical field of three-dimensional perception, and discloses a class-level object pose estimation method based on a learnable prior diffusion model. In the training process, parameters capable of learning prior features can be dynamically updated, a specific category three-dimensional model does not need to be collected, and the category information learning ability of the network is effectively improved. According to the method, the Transform diffusion model is introduced, and the position coding module is combined, so that the understanding capability of the network on the three-dimensional information of the object is enhanced, and the accuracy of object attitude estimation is improved. Aiming at an object with symmetry, the problem of complexity of a plurality of possible attitude solutions caused by symmetry is ingeniously solved from probability distribution of object poses, the interference of symmetry on attitude estimation is effectively avoided, and the estimation performance of a network on the object is remarkably improved. The invention provides an efficient and accurate attitude estimation method, and reliable technical support is provided for grabbing of the intelligent robot.

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, how to accurately perceive three-dimensional objects in the environment has become an important issue in the development of robotics. Object pose estimation is a key task in the robot perception system, which refers to determining the position and pose of an object in three-dimensional space. Only by obtaining accurate object poses can the robot achieve precise manipulation of the object. For example, before grasping a water cup, the robot needs to estimate the direction of the cup mouth to avoid spilling the liquid in the cup during the grasping process.

[0003] Although depth cameras provide robots with depth images and RGB images of objects, there are still multiple challenges in accurately estimating the pose and position of objects from these images. First, there are many types of objects in real scenes, and their sizes and shapes vary significantly. This diversity brings great difficulties to pose estimation. Secondly, some objects are symmetrical, which leads to multiple possible pose solutions for the same observation image, adding complexity to the pose estimation method. Finally, since the depth camera can only obtain the depth image of the object at a certain perspective, it cannot capture the depth information of the occluded part of the object. Therefore, object pose estimation based on incomplete information places higher requirements 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 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, they are based on the premise 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] In patent CN202310963602, a method for adaptive category-level object pose estimation based on graph convolution dual-stream shape prior is reported. This method realizes category-level object pose estimation by learning the three-dimensional model set data of objects of a specific category. Although this method solves the problem that the existing methods can only process specific objects to a certain extent, it still requires the input of a specific category of three-dimensional model set data. This method requires the construction of three-dimensional model data before training, and the network cannot dynamically update the data during the training process, resulting in cumbersome operation steps and limited algorithm learning ability for category information.

[0006] The present invention adopts a Transformer diffusion model regulated by learnable prior features, and provides 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 learns 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 view of the shortcomings of the 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, take their depth images and RGB images, and annotate the pose and size of the target objects. Then, segment and extract the target objects, and construct a dataset with the segmented and extracted images to provide data support for the 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 the DINO v2 module, the PointNet++ module, the position encoding module, the prior feature module, the pose encoding module, the temporal encoding module, and the noise prediction module.

[0014] 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-scale dataset, with a strong ability to understand object semantic information, and is used to improve the processing effect 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 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 K-nearest neighbor algorithms are used to extract the local features of each key point, and the farthest point sampling is used for the sampled point cloud containing local features to halve the number of key points. On this basis, the operations of "extracting local features" and "sampling the farthest point" 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: The input is the 3D position coordinates of 128 key points corresponding to the point cloud features generated by the PointNet++ module, and the 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, so that the position information of the point cloud is embedded into the point cloud feature, improving the model's ability to understand the three-dimensional spatial structure.

[0019] Pose encoding module: extract random pose features through a two-layer MLP network, concatenate the random pose after pose encoding with the point cloud features embedded with position information to generate noise markers.

[0020] 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. This feature updates parameters synchronously with other modules during the training process.

[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 conditional features for noise generation.

[0022] Noise prediction module: It consists of multiple AdaLN modules (adaptive layer normalization modules), multiple attention modules and pose decoding modules. 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.

[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 screens the candidate poses of the 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 candidate poses of the objects, where 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 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, passes through a multi-layer MLP network, perceives the scaling properties of the target object, and finally completes the prediction of the three-dimensional 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 need to be trained in a predetermined order. The pose prediction model should be trained first, then its network parameters should be fixed, and then the pose scoring model and size prediction model should be trained in sequence.

[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 group of depth images and RGB images of the target object obtained in step 1 are 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 square error (MSE) loss function. Specifically, the model uses the Adam optimizer to adjust the parameters in each step of training to reduce the value of the loss function. Repeat the above training process 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.

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

[0031] 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 same training process as the pose prediction model is followed, but the training goal of the pose scoring model is to learn the likelihood probability of the target object pose, rather than directly predicting 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, input it into the pose prediction model, and extract the point cloud features of the target object from its PointNet++ module. Next, 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 loss function for training. Repeat the above training process until the value of the loss function tends to be stable, 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 the training is completed, fix the model parameters and the model can be used for online object pose estimation.

[0036] Whenever the pose of a new similar object needs to be estimated, the depth image and RGB image of the target object are first obtained through the depth camera, and the Mask-RCNN segmentation algorithm is used to segment the image. Subsequently, the segmented image is input into the pose prediction model to extract the point cloud features of the target object. Using the 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 scores of the candidate poses. Next, the candidate poses and the corresponding scores are combined, and the pose aggregation model is used for screening and clustering to output the final pose of the target object. Finally, the point cloud features of the target object are input into the size prediction model to predict the three-dimensional size of the target object.

[0037] The beneficial effects of the present invention are as follows: (1) The present invention does not need to rely on a set of 3D models of objects of the same category, and only requires one training to accurately estimate the pose and size of objects of the same category. (2) The present invention does not require cumbersome special processing of symmetrical objects, and the accuracy of pose estimation for symmetrical objects is significantly improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] Figure 1 It 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 the 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, build a target object dataset according to step 1. Collect the target category objects, the number of objects in each category should be no less than 10, and use a scanner to obtain the 3D model of the target category objects. Then, use a depth camera to record a video sequence of the target object, which should contain about 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 3D model of the target object to calculate the pose of the target object. Finally, segment the target object using the Mask-RCNN segmentation algorithm and extract the image of a single object. If the target category is more common, you can also directly download the 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 in the figure, the specific process is as follows: (1) Constructing a pose prediction model: Figure 2 As shown in the structure 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. Then, 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 pose decoding modules, 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 related models. At this point, the construction of all models is completed.

[0044] Finally, train the object pose estimation model according to the description in step 3. The object pose estimation model contains model parameters and can be trained in three parts: the pose prediction model, the pose scoring model, and the size prediction model, while the pose aggregation model does not need to be trained. The training process must be completed in strict sequence. The specific process is as follows: first train the pose prediction model, then fix the parameters of the model as the starting parameters for training the pose scoring model, and train 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, and train the parameters of the size prediction model.

[0045] During the training process, 95,000 images can be selected as the training set and 5,000 images can be selected as the validation set to evaluate the model training effect. When using the RTX 3090 GPU for training, the batch size is set to 72. After training the pose prediction model for about 72 hours, the loss value of the model can usually converge; next, the batch size is set to 48, and it takes about 24 hours to train the pose scoring model; finally, the batch size is set to 128, and it takes about 24 hours to train the size prediction model. At this point, the training of all models is completed. After the training is completed, the validation set is used to test the running results of the model. If the test results are not ideal, the training time can be appropriately extended until the results meet the requirements. According to experience, when using the RTX 3090 GPU, the complete training task can usually be completed in about 5 days.

[0046] After all models are trained, the position 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, take their depth images and RGB images, and annotate the position and size of the target objects; Subsequently, the target object is segmented and extracted, and a data set is constructed from the segmented and extracted images to provide data support for the 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 poses 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 screens the candidate poses of the 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 candidate poses of the objects, where the category with the largest number of candidate poses of the objects is regarded as 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, passes through a multi-layer MLP network, perceives the scaling properties of the target object, and finally completes the prediction of the three-dimensional size. 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 the training is completed, 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-scale dataset. It has a strong ability to understand object semantic information and is used to improve the processing effect of RGB images. 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 the key point 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; then, multiple groups of MLP networks and K-nearest neighbor algorithms are used to extract the local features of each key point, and the farthest point sampling is used for the sampled point cloud containing local features to halve the number of key points; on this basis, the "extract local features" 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; Position encoding module: The input is the three-dimensional position coordinates of 128 key points corresponding to the point cloud features generated by the PointNet++ module, and the position encoding of the point cloud is output; for any point P in the point cloud, its three-dimensional 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, so as to embed the position information of the point cloud into the point cloud feature, thereby improving the model's ability to understand the three-dimensional spatial structure; Pose encoding module: extract random pose features through a two-layer MLP network, concatenate the random pose after pose encoding with the point cloud features embedded with position information, and generate 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. This feature updates parameters 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: It 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 method for class-level object pose estimation 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 group of depth images and RGB images of the target object obtained in step 1 are randomly selected and input into the pose prediction model; during the training process, random noise is continuously added to the real 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 square error (MSE) loss function; specifically, the model uses the Adam optimizer to adjust the parameters in each step of training to reduce the value of the loss function; repeat the above training process 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; then, the same training process as the pose prediction model is followed, but the training goal of the pose scoring model is to learn the likelihood probability of the target object pose, rather than directly predicting the object pose; (3) Training the size prediction model: The size prediction model is responsible for predicting the length, width, and height of an 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 loss function for training; repeat the above training process until the value of the loss function tends to be stable, indicating that the model training has converged.

4. The method for class-level object pose estimation based on a learnable prior diffusion model according to claim 1, characterized in that: In step 4: whenever the pose of a new similar object needs to be estimated, firstly, the depth image and RGB image of the target object are 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 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