3D对象姿势估计
3D对象姿势估计
对象检测和3D姿态估计在机器人技术中起着至关重要的作用。在各种应用程序中都需要它们,例如导航,对象操纵和检查。Isaac SDK中的3D对象姿势估计应用程序提供了一个框架,可以在模拟中完全训练任何模型的姿势估计,并在模拟以及真实世界中测试和运行推理。
本应用程序中使用的3D姿态估计模型基于Sundermeyer等人的工作。给定RGB图像,该算法首先使用Isaac SDK中可用的任何对象检测模型从具有可用3D CAD模型的已知对象集中检测对象,然后使用自动编码器模型估算其3D姿态。该应用程序旨在通过利用GPU加速同时实现良好的准确性,实现低延迟,实时对象检测和3D姿态估计。
应用概述
../../../_images/3Dpose_estimation_algorithm.png该应用程序包含三个模块:
- 推理模块
- 自动编码器训练模块
- 码本生成模块
推理模块由两个子模块组成:对象检测和3D姿势估计。对象检测子模块拍摄RGB图像,并使用Isaac SDK中可用的任何对象检测模型确定感兴趣对象的边界框。3D姿态估计子模块基于边界框裁剪图像,并估计每个裁剪图像中对象的3D姿态。为了估算对象的3D姿势,此模块需要针对对象的训练有素的编码器神经网络和针对对象的姿势码本。
自动编码器训练模块为对象训练编码器神经网络。该模块使用一种称为增强自动编码器的降噪自动编码器的变体。它训练解码器神经网络来重建输入的RGB图像,以使其对除对象方向以外的任何变化均不变。这些变化包括但不限于背景,照明条件,遮挡,比例,平移等。换句话说,解码器学会了保持对象的方向并规范所有其他方面。
码本生成模块使用训练有素的编码器为对象生成姿态码本。3D姿势估计模块使用此码本在推理时找到与对象方向最佳匹配。
自动编码器训练和代码本生成模块都需要基于用户提供的对象3D CAD模型的模拟数据。该应用程序还使用Unity Engine和Isaac_sim_unity3d提供了模拟数据生成模块,该模块在给定3D CAD模型的情况下为对象生成模拟数据。
以下三个部分将更详细地介绍上述三个模块(自动编码器培训,代码簿生成和推断)中的每个模块。
自动编码器培训
获取和标记训练模型以进行3D姿态估计所需的现实世界数据极具挑战性,耗时且容易出错。对于机器人技术而言,这个问题尤其严重:要求机器人在各种各样的特殊场景中执行姿势估计,因此,对于每种这样的场景,收集大量准确且带有标签的真实世界数据是令人望而却步的,这反过来又减慢了机器人的速度。采用这些模型。
此应用程序使用仿真数据,在仿真过程中完全训练3D姿态估计的自动编码器模型,以通过仿真器中可用的功能(如域随机化和过程生成技术)弥合仿真与实际的差距。
自动编码器培训阶段包括两个主要步骤:
- 生成模拟数据
- 运行自动编码器训练管道
生成模拟数据
训练自动编码器模型需要四种不同类型的数据:
- 渲染的RGB图像
- 对应的分割图像
- RGB图像的去噪版本(即RGB图像的背景,遮挡,照明和平移不变的版本)
- 相应的去噪分割图像
此应用程序允许您通过使用Unity引擎和Isaac_sim_unity3d设置场景,然后将其流式传输到Isaac SDK来生成上述数据。
软件包/ Nvidia / Samples / ObjectPoseEstimation /中的isaac_sim_unity3d 存储库中提供了用于生成上述数据的示例场景。按照IsaacSim Unity3D页面上的说明 在Unity Editor中打开场景。该示例场景可以生成具有随机背景,遮挡对象,光照条件和相机姿势的数据。
场景的重要组成部分如下:
- 彩色摄像机和分割摄像机分别渲染实际的RGB和分割图像。
- 彩色摄像机和分割摄像机分别渲染降噪后的彩色和分割图像。
- 为了使图像不变于背景和遮挡,两台摄像机均使用了剔除蒙版。
- 为了使渲染的彩色图像在光照条件下不变,将自定义反照率明暗器添加到彩色相机中。
- 为了使图像平移不变,将
LookAtFromFixedDistance
脚本添加到两个摄像机,因此它们始终与感兴趣的对象保持恒定的距离。
对场景进行配置,使其可以使用TCP将每帧的所有四个必需标签流式传输到Isaac SDK。场景中默认的关注对象是小车。请按照以下步骤为新对象设置场景:
注意
如果感兴趣的对象已经在packages / Nvidia / Samples / ObjectPoseEstimation / PoseEstimationObjectsGroup中的组中 ,则可以跳过步骤1-3。
- 加载Unity,上传3D CAD模型和对象的纹理,然后创建一个预制件。
- 在预制中,单击“ 添加组件”并添加
LabelSetter
脚本。输入您选择的标签名称。该脚本创建分割图像并计算边界框。 - 将预制件添加到package / Nvidia / Samples / ObjectPoseEstimation中的PoseEstimationObjectsGroup中 。为此,可以将元素尺寸增加1或通过将预制件拖放到列表中来将现有的预制件与对象的预制件交换。
- 将 程序包/Nvidia/Samples/ObjectPoseEstimation/pose_estimation_training.unity场景拖到“ 层次结构”面板中,然后删除面板中的所有现有场景。
../../../_images/class_label_manager.png
-
在“ 类别标签规则”>“小车”>“表达式”字段中,将标签名称设置为预制件。默认名称是“ dolly”,这是“ Dolly”预制件的标签名称。标签名称在isaac_sim_unity3d中用于在场景中渲染对象的分割图像。
-
在“ 类标签”部分中为Dolly元素设置一个名称。该名称用作Isaac SDK中对象的标签。
-
保存场景并播放。
-
在Isaac SDK中,将对象的预制名称添加到packages / object_pose_estimation / apps / training.app.json中组件 的
robot_prefab
参数中 。请注意,这是对象预制的名称,而不是类标签的名称。data.simulation.scenario_manager
例如,如果您将“ TrashCan02.prefab” Unity对象与预制“ trashcan”中的LabelSetter名称一起使用,则
robot_prefab
Isaac应用程序中的参数应为TrashCan02
并且Unity中ClassLabelManager中的表达式名称可以是您选择的任何名称。 。在Isaac应用程序中添加预制名称将在您运行训练应用程序时在场景中生成对象,如以下部分所述。
此时,您应该看到摄像机的位置,对象,照明以及墙壁和地板的背景物质每帧都在变化。您可以根据需要在程序相机****GameObject及其子CameraGroup GameObject中调整背景,相机位置等的随机化频率。
Isaac_sim_unity3d通过TCP套接字与Isaac SDK通信。 isaac / packages / navsim / apps / navsim.app.json使用TcpPublisher节点将模拟数据发布到用户定义的端口。训练应用程序使用TcpSubscriber 接收数据,数据包为package / object_pose_estimation / apps / training.app.json。
运行自动编码器训练管道
使用以下命令运行培训应用程序:
bob@desktop:~/isaac$ bazel run packages/object_pose_estimation/apps:autoencoder_training
注意
确保robot_prefab
在应用程序文件中指定的名称与Unity场景中GameObject的预制名称匹配。
可以在packages / object_pose_estimation / apps / training_config.json文件中设置训练配置 。日志和检查点/tmp/autoenc_logs/ckpts
默认情况下存储在其中,但是可以在训练配置JSON文件中更改此路径。默认情况下,训练应用程序从训练迭代0开始运行,直到迭代次数等于traning_config.json文件中training_step
给出的值为止。
要从中间检查点开始训练,请将checkpoint
config选项设置为相应的检查点文件名,默认情况下,格式为 model-<checkpoint_number>
。该脚本提取附加“ model-”的数字作为检查点编号,然后从该迭代编号重新开始。例如,如果要从迭代编号10000重新启动,请将设置checkpoint
为model-10000
:步骤编号将从10000开始,并以training_steps
config选项的值结束,默认情况下为50000。
要在TensorBoard上查看训练进度,请在终端中运行以下命令:
tensorboard --logdir=/tmp/autoenc_logs/ckpts
可视化界面可从访问http://localhost:6006
。从Isaac_sim_unity3d收到的图像以及边界框可以在的Sight中可视化http://localhost:3000
。
请注意,根据对象的类型,模型可能不需要默认的50000个训练步骤。如果重构的解码器图像(模型的输出)在多个训练样本上清晰可见,则可以更快地结束训练。这是确定模型质量的良好检查。您可以在TensorBoard Images部分中查看重建的图像。
此外,冻结的TensorFlow模型仅在训练结束时生成(即,当迭代次数等于training_steps
config选项的值时)。因此,如果要在完成给定步骤之前结束训练,可以通过运行以下命令从此检查点生成冻结的模型:
bob@desktop:~/isaac$ bazel run packages/ml/tools:freeze_tensorflow_model_tool -- --out /tmp/autoenc_logs/ckpts/model-24000 --output_node_name encoder_output/BiasAdd /tmp/autoenc_logs/ckpts/model-24000
上面的命令将最后一个检查点号设置为24000。您可以根据需要更改此值。
讯息类型
自动编码器训练模块具有以下消息类型:
- ColorCameraProto:保存彩色图像和相机固有信息。
- TensorListProto:定义TensorProto消息的列表,这些消息主要用于传递张量。
- SegmentationCameraProto:保存包含图像中每个像素的类标签的图像。它还包含与相机固有的信息,类似于ColorCameraProto。这用于计算真实的2D边界框。
- RigidBody3GroupProto:包含有关刚体的信息,例如位置,速度,加速度。
- Detections2Proto:保存2D边界框的绝对坐标和类名称。
- Detections3Proto:保留相对于传感器框架及其类别名称检测到的对象的3D姿态。
小码
- TcpSubscriber:由培训应用程序用来从模拟器接收数据。在此示例中,使用了五个TcpSubscriber,每个TcpSubscriber接收编码器和解码器彩色图像及其检测标签,以及模拟中感兴趣的刚体的姿态。
- LabelToBoundingBox:接收SegmentationCameraProto并输出Detections2Proto。该小代码负责根据对象类和实例标签计算地面真理边界框。边界框以1.0的置信度发布为Detections2Proto。
- DetectionImageExtraction:接收包含彩色图像和相机参数的ColorCameraProto,以及包含边界框列表的Detections2Proto。小码裁剪图像中检测到的对象,并将其发布为TensorListProto,其大小为(N xW xH x 3)。
- TensorSynchronization:接受两个或多个TensorListProto输入,并根据其获取时间对其进行同步。此小代码可确保针对每个数据样本同步训练代码接收的所有标签。
- SampleAccumulator:将训练数据标签作为TensorListProto并将其存储在缓冲区中。该小码绑定到Python训练脚本,以便训练脚本可以使用
acquire_samples()
函数直接从该缓冲区采样,该函数将TensorListProto转换为具有相应维的numpy数组列表,并将该列表传递给Python训练脚本。
转换为UFF模型
UFF软件包包含一组实用程序,用于将经过训练的模型从各种框架转换为通用的UFF格式。在此应用程序中,UFF解析器将Tensorflow模型转换为UFF,以便可以将其用于代码本生成和推断。有关更多详细信息,请参阅NVIDIA TensorRT文档。
训练迭代结束时,Tensorflow模型将另存为.pb文件。然后,您需要使用python脚本和UFF解析器将其转换为UFF模型。例如,在24000次迭代结束时,Tensorflow模型另存为, model-24000.pb
并可以使用以下命令转换为UFF模型:
bob@desktop:~/isaac$ bazel run packages/ml/tools:tensorflow_to_tensorrt_tool -- --out /tmp/autoenc_logs/ckpts/ae_model.uff --input_node_name encoder_input --output_node_name encoder_output/BiasAdd /tmp/autoenc_logs/ckpts/model-24000-frozen.pb
码本生成
为了确定推理过程中对象的方向,我们使用了一个包含在不同姿势下对象的潜在空间表示的代码本。将运行时测试图像的潜在向量与码本向量进行比较,以找到最佳匹配。关于密码本生成设置和要求的更多详细信息,可以在Sundermeyer等人的论文中找到。
代码簿生成模块执行以下操作:
- 生成从球体均匀采样的姿势列表,其中感兴趣的对象为中心。
- 将带有采样姿势的远距传送命令发送到模拟器。
- 接收包含每个姿势中对象外观的彩色图像。
- 将包含编码器潜在向量及其相应的姿态和边界框信息的代码簿文件写为JSON文件。
该模块需要UFF模型作为输入,以生成采样姿势的码本。根据您的使用情况,均匀球形采样可以被任何其他采样策略所取代,以实现更好的性能。您还可以通过限制/packages/object_pose_estimation/apps/codebook_generation.app.json应用程序文件中config选项中的俯仰,滚动和偏航角范围来仅采样球体的一部分。
../../../_images/codebook_generation.png在Unity中设置场景
在程序包/ Nvidia / Samples / ObjectPoseEstimation /中的isaac_sim_unity3d存储库中 可以找到在Unity中生成代码本的场景。请按照以下步骤设置场景以生成代码本:
- 加载Unity,将 程序包/Nvidia/Samples/ObjectPoseEstimation/pose_estimation_codebook_generation.unity场景拖到“ 层次结构”面板中,然后删除面板中的所有现有场景。然后将对象预制件拖到场景中,并将其位置设置为原点。
- 在“ 类别标签规则”>“小车”>“表达式”字段中,将标签名称设置为预制件。默认名称是“ dolly”,这是“ Dolly”预制件的标签名称。此标签名称在isaac_sim_unity3d中用于渲染场景中对象的分割图像。
- 在“ 类标签”部分中为Dolly元素设置一个名称。该名称用作Isaac SDK中对象的标签。
注意
此ClassLabelManager
脚本中的所有设置必须与训练场景中相同GameObject的设置匹配。
- 保存场景并播放。
生成密码本
使用在培训结束时保存的UFF模型生成密码本。在Isaac SDK中运行密码本生成应用程序之前,请在/packages/object_pose_estimation/apps/codebook_generation.app.json应用程序文件中配置以下选项 :
- 将对象的预制名称添加到应用文件中组件中的
robot_prefab
选项中simulation.scenario_manager
。默认值为“ Dolly”。 - 将组件中的
whiteList_labels
选项设置为在Unity中“ ClassLabelManager” GameObject的“ 类标签”部分中为对象指定FilterDetectionByLabel
的名称。默认名称是“ Dolly”。在Sight中可视化边界框期间,将使用该名称标记这些框。 -
model_file_path
在TensorRTInference
组件的选项中设置保存的UFF模型的路径。
您还可以codebook_view_sampler
在应用程序文件的组件中设置采样点数,面内滚动,俯仰范围和滚动/偏航角,以从球体中采样姿势。
使用以下命令从终端运行代码簿生成应用程序:
bob @ desktop:〜/ isaac $ bazel运行包/ object_pose_estimation / apps:codebook_generation
此时,在Unity场景中,您应该看到相机位置根据从均匀球体采样的姿势改变每一帧,这是默认选项。渲染的对象可以在的Sight中可视化http://localhost:3000
。默认情况下,密码簿保存在 /tmp/codebook.jsonl中。
小码
自动编码器训练模块中使用的许多小码在代码本生成模块中重新使用。下面列出了使用的其他小码:
- CodebookSampler:假设要观察的对象位于原点,并为每个刻度发布一个采样的姿势,以采样照相机姿势以生成代码簿。它使用二十面体细分对来自给定半径球体的视点进行均匀采样,并沿相机轴附加平面内旋转(滚动)。
- TcpPulisher:将数据从Isaac SDK发送到模拟器。对于生成码本,使用一个TcpPublisher发送相机应在模拟中传送到的姿势。
- Detections3Encoder:接收包含姿势信息的Detections3Proto,并将其编码为张量。张量以TensorListProto的形式发布,其中包含所有检测到的对象的姿态信息。
- RigidBodyToDetections3:接收一个RigidBody3GroupProto,其中包含在Isaac SDK坐标框中具有姿势的3D刚体列表,如果需要,将其转换为参考框作为输入刚体之一,并在该3D刚体中发布该列表。参考框架为Detections3Proto。该小码用于计算从模拟获得的相对于相机框架的刚体的姿态。结果输出姿势用于训练姿势估计模型。
- TensorRTInference:将冻结的神经网络模型加载到内存中,生成优化的TensorRT引擎,使用网络输入作为TensorListProto评估模型,并发布网络输出,该输出是对象类所有实例的嵌入矢量,如下所示:一个TensorListProto。
- ImagePoseEncoder:获取ColorCameraProto的照相机固有属性,包含边界框列表的Detections2Proto以及RigidBody3GroupProto以获得刚体的3D姿势。它将从这些输入中提取的七个标签发布为TensorListProto中的单个张量,以生成代码本。
- CodebookWriter:接收两条TensorListProto消息,一条消息包含嵌入列表,另一条消息包含从ImagePoseEncoder小码接收到的相应解码信息,并将其作为JSON Proto发送。在此示例中,准备了这样的代码本,其中代码是来自编码器的潜在矢量,而解码信息是这些矢量的姿态以及边界框和相机参数。
- JsonWriter:接收JsonProto消息并将其写为JSON文件。此小代码从CodebookWriter小代码收集姿势编码,并将其写入完成的代码本文件中,以供推理模块使用。
推理
从单个RGB图像进行的姿势估计模型的端到端推断包括两个主要子图。
- 对象检测子图,可以利用Isaac SDK中的任何对象检测模型
- 姿势估计子图,该子图运行编码器网络,并根据来自生成的代码本的潜矢量的最佳匹配来估计3D姿势。推理应用程序为每个对象类实例化一个姿势估计子图。可以在单个姿势估计子图中同时处理一个对象类的多个实例。
具有不同CAD模型的每个对象都需要相应的训练模型和生成的密码本。然后,必须为每个对象向推理应用程序添加一个单独的子图。场景中同一对象的不同实例将自动传递到与该对象类相对应的单个子图。
运行推断
根据数据收集方式有四种推理应用程序:对作为模拟模型准确性的第一个测试的推理数据进行推理,使用摄像头馈送进行实时推理,使用示例图像进行推理以及使用记录的日志进行推理。
首先,对推理应用/packages/object_pose_estimation/apps/codebook_generation.app.json
程序文件执行与“ 密码本生成”部分中针对应用程序文件指示的相同配置步骤 。推理应用程序文件位于 / packages / object_pose_estimation / apps中。
- 要在模拟数据上运行推理应用程序,请先在Unity中播放姿势估计训练场景,然后在Isaac SDK中运行以下命令:
bob @ desktop:〜/ isaac $ bazel运行包/ object_pose_estimation / apps:pose_estimation_inference_sim
- 要使用实时摄像机供稿运行推理应用程序,请在Isaac SDK中运行以下命令:
bob @ desktop:〜/ isaac $ bazel运行包/ object_pose_estimation / apps:pose_estimation_inference_camerafeed
- 要使用示例图像运行推理应用程序,请在Isaac SDK中运行以下命令:
bob @ desktop:〜/ isaac $ bazel运行包/ object_pose_estimation / apps:pose_estimation_inference_imagefeeder
- 要使用记录的日志运行推理应用程序,请在Isaac SDK中运行以下命令:
bob @ desktop:〜/ isaac $ bazel运行包/ object_pose_estimation / apps:pose_estimation_inference_replay
推理应用程序从不同的输入源获取RGB图像,并输出检测到的对象的估计3D姿势。提供了手推车对象的样本数据,以使用图像和日志来运行最后两个推理应用程序。估计的姿势可以在的Sight中自动显示http://localhost:3000
。
估计姿势有两种可视化类型:
- 一个3D边界框,需要指定零方向的3D边界框大小以及从对象中心到边界框中心的转换。
ObjectDetectionViewer
在推理应用程序文件的组件中配置这些参数。 - 场景中CAD模型的渲染,它要求对象CAD模型和文件名的路径。这些分别对应于推理应用程序文件中组件中的
assetroot
和assets
参数websight
。
小码
- 实感:从实感摄像头启用数据摄取。将原始图像编码为包含RGB图像和相机固有信息的ColorCameraProto消息。
- FilterDetectionsByLabel:接收边界框检测列表作为Detections2Proto,过滤所需对象类的检测,并将其发布为Detections2Proto。
- DetectionImageExtraction:接收包含彩色图像和相机参数的ColorCameraProto以及包含边界框列表的Detections2Proto。小码裁剪图像中检测到的对象,并将其发布为TensorListProto,其大小为(N xW xH x 3)。
- TensorRTInference:将冻结的神经网络模型加载到内存中,生成优化的TensorRT引擎,使用网络输入作为TensorListProto评估模型,然后将网络输出(即对象类所有实例的嵌入矢量)发布为TensorListProto。 。
- CodebookLookup:获取包含嵌入的TensorListProto,并使用余弦相似度找到向量的前K个最佳匹配列表与Codebook中的代码列表。它将所有实例的最佳匹配嵌入向量的相应标签发布为TensorListProto。
- PoseEstimation:获取TensorListProto,其中包含来自代码簿的最佳匹配矢量的标签,并将对象的姿态估计发布为Detections3Proto。