如何使用OpenUSD构建支持生成式AI的合成数据工作流

训练用于驱动自主机器（例如机器人和自动驾驶汽车）的物理 AI 模型需要大量数据。获取大量多样化的训练数据可能很困难、耗时且昂贵。数据通常因隐私限制或担忧而受到限制，或者对于新颖的用例来说可能根本不存在。此外，可用数据可能不适用于各种潜在情况的全部范围，从而限制了模型准确预测和响应各种场景的能力。

合成数据，通过计算机模拟从数字孪生中生成，可替代真实数据，使开发者能够启动物理 AI 模型训练。您可以快速生成大量多样化的数据集，通过改变许多不同的参数，如布局、资产放置、位置、颜色、物体大小和光照条件。然后，利用这些数据帮助创建通用模型。

实现逼真度对缩小模拟与现实领域的差距至关重要。该过程旨在使用正确的属性（例如材质和纹理）表示虚拟环境中的每个对象，以准确地模仿其真实世界的表示。如果没有 AI 的帮助，这是一个耗时的手动过程。生成式 AI 可以帮助加速从资产创建到代码生成等过程的许多方面，支持开发者构建稳健且多样化的训练数据集。

本文介绍了如何使用适用于通用场景描述的 NVIDIA NIM 微服务以及 NVIDIA Omniverse Replicator 构建自定义合成数据生成（SDG）管线。NVIDIA NIM 是一套加速推理微服务，允许组织在任何位置——云、数据中心、工作站和 PC 上——在 NVIDIA GPU上运行 AI 模型。Omniverse Replicator 是基于通用场景描述（OpenUSD）和 NVIDIA RTX 构建的 SDK。

本文还将介绍如何获取生成的图像，并使用示例 ComfyUI 工作流程作为参考管线的一部分进一步增强这些图像。然后，生成的图像可用于预训练模型和工具，例如 TAO、PyTorch 或 TensorFlow。

参考工作流程概述

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

图 1. NVIDIA Isaac Sim 中的机器人基础模型端到端参考工作流

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

生成带注释的合成数据后，我们将执行生成式 AI 增强，使用简单的提示进一步分散和扩展数据集。

使用 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 函数，在给定的 stage 中，在 x 和 y 轴上的 6×5 网格中添加矩形 lights，尺寸为 x：5 和 y：10。每个矩形 lights 的强度应为 5000，曝光值应为 1，大小为 (10,3)。lights 应在 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 个新 lights (图 2)。此过程可用于添加更多资产 (如 shelves、boxes 和 forklifts)，以完成场景。

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

如果需要其他资产或背景来增强场景，还可以使用使用 NVIDIA Edify 构建的服务。NVIDIA Edify 是一种功能强大的多模态架构，用于构建 AI 模型，生成视觉内容，例如 4K 图像、详细的 3D 网格、16K 360 HDRi、PBR 材质和视频。然后，我们会对 AI 模型进行优化和打包，以通过 NVIDIA NIM 实现最大性能。这将加速内容创作过程。

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

Shutterstock Generative 3D API 已进入商业测试阶段，可通过 TurboSquid by Shutterstock 进行访问。

此外，fVDB 是一个开源深度学习框架，可用于生成大规模场景，利用真实世界的 3D 数据训练空间智能。它基于 OpenVDB 构建 AI 运算符，以创建真实环境的高保真虚拟表示，包括神经辐射场神经辐射场(NeRF)、点云表面重建，甚至大规模生成式 AI。这些丰富的 3D 数据集都是 AI-ready 的，支持高效的模型训练和推理。

生成域随机化代码

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

但是，为大规模域随机化编写代码可能会很繁琐，并且会拖慢数据生成的迭代过程。解决方案是什么？利用 USD Code NIM 的强大功能充当 copilot。

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

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

首先，输入以下提示：

在开放舞台中，我想使用一个函数，仅在 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 作为 copilot 会比自己编写代码更快地实现无差错代码。

最终代码如下所示：

 

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 编写以支持开发可在不同场景中轻松重复使用的辅助函数，从而提高未来运行的效率。

导出带注释的图像

通过在 Replicator 中设置域随机化，就可以导出第一批带注释的图像。Replicator 具有许多开箱即用的注释器，例如二维边界框、语义分割、深度、法线等。输出类型 (例如边界框或分割) 取决于模型类型或用例。数据可以使用 BasicWriter 输出为基本数据，使用 KittiWriter 输出为 KITTI 数据，或使用自定义写入器输出为 COCO 数据。

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

图 6. Replicator 可以生成各种带注释的数据，包括(从左上角顺时针开始) Normals, RGB, Depth, 和 Semantic Segmentation。

使用 ComfyUI 增强合成数据集

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

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

有关 ComfyUI 图形的入门参考，请参阅Generative AI for Digital Twins Guide：https://docs.omniverse.nvidia.com/genai-nim-digitaltwin.html

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

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

启动‘queue prompt’后，在图 7 的最右侧可以看到增强输出。这些输出是传统渲染的合成数据和增强区域的组合。模型通常能够从更广泛的图像中获取照明线索和上下文，并对增强区域进行适当的照明或阴影处理。

地板类型和物体颜色等细节可以在现有图像中进行更改。以下四个提示的结果图像如图 8 所示。

提示 1

白色平铺油毡地板
绿色闪亮的新平衡重叉车

木托盘浅色松木、软木
垃圾容器

提示 2

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

提示 3

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

提示 4

绿色切口油毡地板
蓝色生锈平衡重叉车
木托盘浅色松木、软木
垃圾容器

 

训练模型

虽然本文并未明确介绍，但下一步逻辑上应该是训练计算机视觉模型。您可以启动 NVIDIA 预训练计算机视觉模型或选择您自己的模型。然后，可以使用 NVIDIA TAO 等训练框架对预训练模型进行微调。TAO 基于 TensorFlow 和 PyTorch 构建，并使用迁移学习来加速模型训练过程。与使用真实数据一样，您可能必须经过多次迭代才能对模型 KPI 感到满意。

如果您已经设置了管线，就可以返回 3D 仿真环境，通过更改附加参数和使用 ComfyUI 工作流增强参数来生成新数据。自动化管线减少了生成和标注新数据所需的时间，而这通常是训练模型时的瓶颈。

总结

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

我们很期待看到您如何使用这一工作流程来开发自己的 SDG 管道。要开始使用，请查看详细的端到端工作流。