基于OpenUSD的合成数据生成

DATASET Dec 1, 2024

训练用于驱动机器人和自动驾驶汽车等自主机器的物理 AI 模型需要大量数据。获取大量不同的训练数据可能很困难、耗时且成本高昂。由于隐私限制或担忧,数据通常有限,或者可能根本不存在于新用例中。此外,可用数据可能不适用于所有潜在情况,从而限制了模型准确预测和响应各种场景的能力。

通过计算机模拟从数字孪生生成的合成数据为现实世界数据提供了一种替代方案,使开发人员能够引导物理 AI 模型训练。您可以通过改变许多不同的参数(例如布局、资产放置、位置、颜色、对象大小和照明条件)来快速生成大量不同的数据集。然后可以使用这些数据来帮助创建通用模型。

实现照片级真实感对于缩小模拟与现实领域的差距至关重要。此过程旨在使用正确的属性(例如材质和纹理)表示虚拟环境中的每个对象,以准确模仿它们在现实世界中的表示。如果没有 AI 的帮助,这是一个手动且耗时的过程。生成式 AI 可以帮助加快该过程的许多方面,从资产创建到代码生成,支持开发人员构建强大而多样化的训练数据集。

本文介绍了如何使用 NVIDIA NIM 微服务NVIDIA Omniverse Replicator为 USD 构建自定义合成数据生成 (SDG) 管道。NVIDIA NIM 是一组加速推理微服务,允许组织在任何地方(云端、数据中心、工作站和 PC)在 NVIDIA GPU 上运行 AI 模型。Omniverse Replicator 是基于通用场景描述 (OpenUSD) 和 NVIDIA RTX 构建的 SDK。

本文还将向你展示如何获取生成的图像并使用示例 ComfyUI 工作流作为参考管道的一部分进一步增强它们。然后可以将生成的图像与预训练模型和工具(如 TAO、PyTorch 或 TensorFlow)一起使用。

1、参考工作流程概述

工作流程从现有仓库的 3D 场景开始,其中包含所有必要的 3D 资产,例如货架、箱子、托盘等。要了解有关创建工厂数字孪生的更多信息,请参阅工作流程示例。为了进一步增强 3D 场景,可以使用 3D NIM 微服务添加更多资产,并更改 360 HDRI 背景图像以实现更多随机化。

图 1. 使用 NVIDIA Isaac Sim 训练机器人基础模型的端到端参考工作流程

下一步涉及使用 USD Code NIM 生成域随机化所需的代码,USD Code NIM 是一种最先进的大型语言模型 (LLM),可回答 OpenUSD 知识查询并生成 USD-Python 代码。域随机化是创建合成数据的重要且必要的步骤。场景中的许多参数都可以在 Omniverse Replicator 中以编程方式更改。

生成带标注的合成数据后,我们将使用简单的提示执行生成式 AI 增强,以进一步多样化和扩展数据集。

2、使用 USD NIM 微服务生成场景

USD Code 使你能够在现有的 3D 开发工作流程中更高效地学习和开发 OpenUSD,只需输入提示并获得返回响应即可。USD Code 可以回答你的问题以及编写自定义 OpenUSD Python 代码。

你可以在 NVIDIA NIM API 目录中或通过 Visual Studio Code 尝试 USD Code,并在 Omniverse Kit 中将其用作示例扩展来修改 OpenUSD 阶段。有关详细信息,请参阅 USD Code NIM 文档。

请参阅下面的生成代码示例提示:

编写一个 Python 函数,给定一个阶段,在 x 和 y 轴上的 6×5 网格中添加矩形灯,尺寸为 x:5 和 y:10。每个矩形光的强度应为 5000,曝光度应为 1,大小为 (10,3)。光应在 z 轴上对齐,并嵌套在名为 Lights 的 xform 下。

生成的代码如下所示:

from pxr import Sdf, Usd, UsdGeom, UsdLux
 
def create_rect_lights(stage: Usd.Stage) -> None:
    """Create a 6x5 grid of rect lights on the x and y axis."""
    lights_xform = UsdGeom.Xform.Define(stage, Sdf.Path("/Lights"))
     
    for i in range(6):
        for j in range(5):
            light_path = Sdf.Path(f"/Lights/Light_{i}_{j}")
            rect_light = UsdLux.RectLight.Define(stage, light_path)
            rect_light.CreateIntensityAttr().Set(5000)
            rect_light.CreateExposureAttr().Set(1)
            rect_light.CreateSizeAttr().Set((10, 3))
            rect_light.AddTranslateOp().Set((i * 5, j * 10, 0))
            lights_xform.AddChild(rect_light.GetPrim())

结果代码在执行时会按照规定的间距在仓库中创建并填充 30 个新灯(图 2)。此过程可用于添加更多资产(如货架、箱子和叉车)以完成场景。

图 2. 使用 USD Code NIM 从代码生成的仓库场景

如果你需要额外的资产或背景来增强场景,还可以使用使用 NVIDIA Edify 构建的服务,这是一种强大的多模式架构,用于构建用于生成视觉内容的 AI 模型,例如 4K 图像、详细的 3D 网格、16K 360 HDRi、PBR 材料和视频。然后使用 NVIDIA NIM 对 AI 模型进行优化和打包,以实现最佳性能。这加快了内容创建过程。

借助 Shutterstock 提供的由 NVIDIA Edify 提供支持的生成式 3D,你可以在 10 秒内使用文本提示或参考图像生成网格预览。然后,你可以在几分钟内使用 PBR 材质生成可编辑的网格,从而实现快速布景、概念设计或原型设计。此外,360 HDRi 生成功能也由 NVIDIA Edify 提供支持,使用户能够使用文本或图像提示生成 16K 360 HDRi,以生成背景并匹配 3D 场景的灯光。

Shutterstock 生成式 3D API 处于商业测试阶段,可通过 Shutterstock 的 TurboSquid 访问。

此外,fVDB 是一个开源深度学习框架,可用于生成大规模场景,以使用真实世界的 3D 数据训练空间智能。它在 OpenVDB 之上构建 AI 运算符,以创建真实世界环境的高保真虚拟表示,包括神经辐射场 (NeRF)、点云的表面重建,甚至大规模生成式 AI。这些丰富的 3D 数据集已为 AI 做好准备,可进行高效的模型训练和推理。要了解更多信息,请参阅使用深度学习框架 fVDB 从现实世界 3D 数据构建空间智能

3、生成域随机化代码

域随机化是向数据集添加多样性的重要技术。它是 Omniverse Replicator 的核心功能之一。你可以以编程方式更改给定场景中的任意数量的变量,包括照明、对象位置、材质和纹理等。创建多样化的数据集有助于感知模型在许多不同场景中执行。

但是,编写大规模域随机化的代码可能很繁琐,并且会减慢数据生成的迭代过程。解决方案?利用 USD Code NIM 的强大功能充当副驾驶。

本节将引导你了解如何使用 USD Code NIM 生成代码作为域随机化的起点。你可以选择在 Visual Studio Code 中激活 USD Code,也可以直接在任何基于 Omniverse Kit 的应用程序中激活,其中 ChatUSD 代理扩展可以打开 USD Code 窗口以输入提示。

图 3. ChatUSD 代理扩展可以在 Omniverse 中激活

首先,输入以下提示:

在打开阶段,我希望有一项功能,可以仅在 x 轴上将名为“/Root/RectLight_03”的灯光在 0 到 -20 米之间随机移动。
import omni.replicator.core as rep
from omni.replicator.core.distribution import uniform
import asyncio
import omni.kit
 
stage = omni.usd.get_context().get_stage()
camera = '/OmniverseKit_Persp'
render_product = rep.create.render_product(camera, (1024, 1024))
 
def move_light() -> None:
    pass
 
# Initialize and attach writer
writer = rep.writers.get("BasicWriter")
writer.initialize(output_dir="_output", rgb=True)
writer.attach([render_product])
 
async def go(num_frames=10):
    for _ in range(num_frames):
         move_light()
         await rep.orchestrator.step_async()
 
asyncio.ensure_future(go())”
图 4. USD Code NIM 对用灯光填充仓库的初始响应

你可以通过限制灯光的 Y 和 Z 位置、检查灯光是否具有正确的变换等方式进一步改进这一点。请注意,虽然这是一个迭代过程,但使用 USD Code 作为副驾驶将帮助您比自己编写代码更快地获得无错误代码。

最终代码看起来会像这样:

import omni.replicator.core as rep
from omni.replicator.core.distribution import uniform
from pxr import Usd, UsdGeom, Gf
import asyncio
import omni.kit
 
stage = omni.usd.get_context().get_stage()
camera = '/OmniverseKit_Persp'
render_product = rep.create.render_product(camera, (1024, 1024))
 
import random
 
def move_light() -> None:
    """Randomly move the light named "/Root/RectLight_03" between 0 and -20 meters on only the x-axis."""
    light_prim = stage.GetPrimAtPath("/Root/RectLight_03")
    translate_attr = light_prim.GetAttribute("xformOp:translate")
    if translate_attr:
        current_translation = translate_attr.Get()
        new_x = random.uniform(-20, 0)  # random x value between -20 and 0
        new_translation = Gf.Vec3d(new_x, current_translation[1], current_translation[2])
        translate_attr.Set(new_translation)
    else:
        new_x = random.uniform(-20, 0)  # random x value between -20 and 0
        light_prim.AddAttribute("xformOp:translate", Sdf.ValueTypeNames.Float3).Set(Gf.Vec3d(new_x, 0, 0))
 
# Initialize and attach writer
writer = rep.writers.get("BasicWriter")
writer.initialize(output_dir="_output", rgb=True, normals=True, distance_to_image_plane=True, semantic_segmentation=True)
writer.attach([render_product])
 
async def go(num_frames=10):
    for _ in range(num_frames):
         move_light()
         await rep.orchestrator.step_async()
 
asyncio.ensure_future(go())

在 Replicator 脚本编辑器中执行代码将导致仓库中的灯光随机打开和关闭(图 5)。

图 5. 作为场景随机化的一部分,仓库灯打开和关闭,由 Replicator 代码控制

此示例代码只是使用 USD Code NIM 进行域随机化的众多示例之一。您可以通过向场景添加随机化并增加数据集的多样性来继续迭代。您还可以指定您的 Python 编写以支持开发可轻松重用于不同场景的辅助函数。这提高了未来运行的效率。

4、导出带标注的图像

在 Replicator 中设置域随机化后,你现在可以导出第一批带标注的图像。Replicator 有许多开箱即用的标注器,例如 2D 边界框、语义分割、深度、法线等等。输出的类型(例如边界框或分割)将取决于模型或用例的类型。可以使用 BasicWriter 将数据导出为基本数据,使用 KittiWriter 将数据导出为 KITTI,或使用自定义编写器将数据导出为 COCO。

更重要的是,从 Replicator 生成的数据捕获各种物理交互,如刚体动力学(例如运动和碰撞)或光在环境中的交互方式。图 6 显示了可以从 Replicator 导出的带注释数据类型的示例。

图 6. Replicator 可以生成各种带标注的数据,包括(从左上角开始顺时针方向)法线、RGB、深度和语义分割

5、使用 ComfyUI 增强合成数据集

进一步增强合成数据集可以创建新的变化,例如更改背景并添加额外的纹理和材质细节,所有这些都使用文本提示。广泛的变化不仅有益,而且可以通过高度逼真的结果缩小外观域差距。总的来说,花在数据集上的时间、大量高质量数字资产的负担以及在增强阶段重新生成新变化的灵活性可以节省大量时间。

ComfyUI 是一个基于 Web 的后端和 GUI,用于使用扩散模型构建和执行管道。它是一个功能强大的开源工具,可在 GitHub 上下载。它可以与 SDXL、其他扩散模型或你选择的微调模型一起使用。更广泛的社区已经创建了一个额外的技术和功能生态系统,作为 ComfyUI 的节点,你可以将其集成到你的图中。

有关入门的参考 ComfyUI 图表,请参阅《数字孪生指南的生成式 AI》

图 7. 使用扩散模型增强合成生成的图像的示例 ComfyUI 工作流程

从高层次来看,可以将图视为“使用模型编程”。区域提示用于指导从数据集图像生成的扩散输出。ControlNet 节点接收从法线和分割创建的轮廓图像。ControlNets 与区域提示相结合,是控制变化的关键因素,同时保留了此数据集中一致性的重要结构。

在启动“队列提示”后,可以在图 7 的最右侧看到增强的输出。输出是传统渲染的合成数据和增强区域的组合。模型通常能够从更广泛的图像中获取照明线索和背景,并适当地照亮或阴影增强区域。

可以在现有图像中更改地板类型和物体颜色等细节。四个提示

如下所示,结果图像如图 8 所示。

提示 1:

白色瓷砖油毡地板
绿色闪亮新平衡叉车
木托盘浅色松木、软木
垃圾箱

提示 2:

深色开裂脏混凝土地板
黄色平衡叉车
木托盘浅色松木、软木
黑色垃圾箱

提示 3:

开裂混凝土地板
白色平衡叉车
木托盘浅色松木、软木
垃圾箱

提示 4:

绿色碎裂油毡地板
蓝色生锈平衡叉车
木托盘浅色松木、软木
垃圾箱
图 8. 合成生成的图像,显​​示不同类型的地板和叉车颜色

6、训练模型

虽然这篇文章没有明确介绍,但合乎逻辑的下一步是训练计算机视觉模型。您可以启动 NVIDIA 预训练的计算机视觉模型或选择自己的模型。然后可以使用 NVIDIA TAO 等训练框架对预训练模型进行微调。TAO 基于 TensorFlow 和 PyTorch 构建,并使用迁移学习来加速模型训练过程。就像处理真实数据一样,你可能需要进行几次迭代才能对模型 KPI 感到满意。

假设你已经设置了管道,你可以返回 3D 模拟环境,通过更改其他参数并使用 ComfyUI 工作流对其进行扩充来生成新数据。自动化管道减少了生成和标记新数据所需的时间,这通常是训练模型时的瓶颈。

7、结束语

你可以使用 NVIDIA NIM 微服务和 NVIDIA Omniverse Replicator 快速轻松地构建自定义合成数据生成管道,如本文所述。你可以使用 ComfyUI 进一步增强生成的图像。然后可以将这些生成的图像与预训练模型和工具(例如 NVIDIA TAO、PyTorch 或 TensorFlow)一起使用。

我们很高兴看到你如何使用此工作流程开发自己的 SDG 管道。要开始使用,请查看详细的端到端工作流程。

在 SIGGRAPH 2024 上,NVIDIA 首席执行官 Jensen Huang 与 Meta 创始人兼首席执行官 Mark Zuckerberg 和 WIRED 高级作家 Lauren Goode 进行了炉边谈话。可按需观看 NVIDIA 在 SIGGRAPH 2024 上的炉边谈话和其他会议。


原文链接:How to Build a Generative AI-Enabled Synthetic Data Pipeline with OpenUSD

汇智网翻译整理,转载请标明出处

Tags