如何使用OpenUSD構建支持生成式AI的合成數(shù)據(jù)工作流

訓練用于驅動自主機器（例如機器人和自動駕駛汽車）的物理 AI 模型需要大量數(shù)據(jù)。獲取大量多樣化的訓練數(shù)據(jù)可能很困難、耗時且昂貴。數(shù)據(jù)通常因隱私限制或擔憂而受到限制，或者對于新穎的用例來說可能根本不存在。此外，可用數(shù)據(jù)可能不適用于各種潛在情況的全部范圍，從而限制了模型準確預測和響應各種場景的能力。

合成數(shù)據(jù)，通過計算機模擬從數(shù)字孿生中生成，可替代真實數(shù)據(jù)，使開發(fā)者能夠啟動物理 AI 模型訓練。您可以快速生成大量多樣化的數(shù)據(jù)集，通過改變許多不同的參數(shù)，如布局、資產(chǎn)放置、位置、顏色、物體大小和光照條件。然后，利用這些數(shù)據(jù)幫助創(chuàng)建通用模型。

實現(xiàn)逼真度對縮小模擬與現(xiàn)實領域的差距至關重要。該過程旨在使用正確的屬性（例如材質和紋理）表示虛擬環(huán)境中的每個對象，以準確地模仿其真實世界的表示。如果沒有 AI 的幫助，這是一個耗時的手動過程。生成式 AI 可以幫助加速從資產(chǎn)創(chuàng)建到代碼生成等過程的許多方面，支持開發(fā)者構建穩(wěn)健且多樣化的訓練數(shù)據(jù)集。

本文介紹了如何使用適用于通用場景描述的 NVIDIA NIM 微服務以及 NVIDIA Omniverse Replicator 構建自定義合成數(shù)據(jù)生成（SDG）管線。NVIDIA NIM 是一套加速推理微服務，允許組織在任何位置——云、數(shù)據(jù)中心、工作站和 PC 上——在 NVIDIA GPU上運行 AI 模型。Omniverse Replicator 是基于通用場景描述（OpenUSD）和 NVIDIA RTX 構建的 SDK。

本文還將介紹如何獲取生成的圖像，并使用示例 ComfyUI 工作流程作為參考管線的一部分進一步增強這些圖像。然后，生成的圖像可用于預訓練模型和工具，例如 TAO、PyTorch 或 TensorFlow。

參考工作流程概述

該工作流程從現(xiàn)有倉庫的 3D 場景開始，該場景包含所有必要的 3D 資產(chǎn)，例如貨架、箱子、托盤等。如需了解有關創(chuàng)建工廠數(shù)字孿生的更多信息，請參閱工作流程示例。為了進一步增強 3D 場景，可以使用 3D NIM 微服務添加更多資產(chǎn)，并更改 360 HDRI 背景圖像以實現(xiàn)更多隨機化。

圖 1. NVIDIA Isaac Sim 中的機器人基礎模型端到端參考工作流

下一步是使用 USD Code NIM 生成域隨機化所需的代碼，USD Code NIM 是一種先進的大型語言模型(LLM)，可回答 OpenUSD 知識查詢并生成 USD-Python 代碼。域隨機化是創(chuàng)建合成數(shù)據(jù)的強大且必要的步驟。在 Omniverse Replicator 中，許多場景參數(shù)可以以編程方式進行更改。

生成帶注釋的合成數(shù)據(jù)后，我們將執(zhí)行生成式 AI 增強，使用簡單的提示進一步分散和擴展數(shù)據(jù)集。

使用 USD NIM 微服務生成場景

USD Code 使您能夠在現(xiàn)有的 3D 開發(fā)工作流程中更高效地學習和開發(fā) OpenUSD，只需通過輸入提示并獲得回復。USD Code 既可以回答問題，也可以編寫自定義的 OpenUSD Python 代碼。

您可以在 NVIDIA NIM API 目錄中或通過 Visual Studio Code 試用 USD Code，在 Omniverse Kit 中將其用作修改 OpenUSD 階段的示例擴展程序。有關詳細信息，請參閱USD Code NIM 文檔。

請參閱下面的示例提示以生成代碼：

編寫一個 Python 函數(shù)，在給定的 stage 中，在 x 和 y 軸上的 6×5 網(wǎng)格中添加矩形 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())

執(zhí)行后生成的代碼會在倉庫中按照規(guī)定的間距創(chuàng)建并填充 30 個新 lights (圖 2)。此過程可用于添加更多資產(chǎn) (如 shelves、boxes 和 forklifts)，以完成場景。

圖 2. 使用 USD Code NIM 從代碼生成的倉庫場景

如果需要其他資產(chǎn)或背景來增強場景，還可以使用使用 NVIDIA Edify 構建的服務。NVIDIA Edify 是一種功能強大的多模態(tài)架構，用于構建 AI 模型，生成視覺內容，例如 4K 圖像、詳細的 3D 網(wǎng)格、16K 360 HDRi、PBR 材質和視頻。然后，我們會對 AI 模型進行優(yōu)化和打包，以通過 NVIDIA NIM 實現(xiàn)最大性能。這將加速內容創(chuàng)作過程。

借助 NVIDIA Edify 提供支持的 Shutterstock Generative 3D，您可以在 10 秒內通過文本提示或參考圖像生成網(wǎng)格預覽。然后在幾分鐘內使用 PBR 材質生成可隨時編輯的網(wǎng)格，從而實現(xiàn)快速的布景、概念或原型設計。此外，360 HDRi 生成功能也由 NVIDIA Edify 提供支持，使用戶能夠使用文本或圖像提示生成 16K 360 HDRi，以生成背景并匹配 3D 場景的照明。

Shutterstock Generative 3D API 已進入商業(yè)測試階段，可通過 TurboSquid by Shutterstock 進行訪問。

此外，fVDB 是一個開源深度學習框架，可用于生成大規(guī)模場景，利用真實世界的 3D 數(shù)據(jù)訓練空間智能。它基于 OpenVDB 構建 AI 運算符，以創(chuàng)建真實環(huán)境的高保真虛擬表示，包括神經(jīng)輻射場神經(jīng)輻射場(NeRF)、點云表面重建，甚至大規(guī)模生成式 AI。這些豐富的 3D 數(shù)據(jù)集都是 AI-ready 的，支持高效的模型訓練和推理。

生成域隨機化代碼

域隨機化是向數(shù)據(jù)集添加多樣性的一項重要技術。它是 Omniverse Replicator 的核心功能之一。您可以通過編程更改給定場景中的任意數(shù)量的變量，包括照明、對象位置、材質和紋理等。創(chuàng)建多樣化的數(shù)據(jù)集有助于感知模型在許多不同的場景中執(zhí)行操作。

但是，為大規(guī)模域隨機化編寫代碼可能會很繁瑣，并且會拖慢數(shù)據(jù)生成的迭代過程。解決方案是什么？利用 USD Code NIM 的強大功能充當 copilot。

本節(jié)介紹如何使用 USD Code NIM 生成代碼，以此作為域隨機化的起點。您可以選擇在 Visual Studio Code 中激活 USD Code，也可以直接在任何基于 Omniverse Kit 的應用程序中激活 USD Code，其中 ChatUSD 代理擴展程序可以打開 USD Code 窗口，以便輸入提示。

圖 3. Omniverse 中可以激活 ChatUSD 代理擴展程序

首先，輸入以下提示：

在開放舞臺中，我想使用一個函數(shù)，僅在 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 會比自己編寫代碼更快地實現(xiàn)無差錯代碼。

最終代碼如下所示：

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 腳本編輯器中執(zhí)行代碼將實現(xiàn)倉庫中的燈光隨機地打開和關閉(圖 5)。

圖 5. 作為場景隨機化的一部分，倉庫燈在 Replicator 代碼的指導下開啟和關閉

此示例代碼只是使用 USD Code NIM進行域隨機化的眾多示例之一。您可以通過在場景中添加隨機化并增加數(shù)據(jù)集的多樣性來繼續(xù)迭代。您還可以指定 Python 編寫以支持開發(fā)可在不同場景中輕松重復使用的輔助函數(shù)，從而提高未來運行的效率。

導出帶注釋的圖像

通過在 Replicator 中設置域隨機化，就可以導出第一批帶注釋的圖像。Replicator 具有許多開箱即用的注釋器，例如二維邊界框、語義分割、深度、法線等。輸出類型 (例如邊界框或分割) 取決于模型類型或用例。數(shù)據(jù)可以使用 BasicWriter 輸出為基本數(shù)據(jù)，使用 KittiWriter 輸出為 KITTI 數(shù)據(jù)，或使用自定義寫入器輸出為 COCO 數(shù)據(jù)。

更重要的是，Replicator 生成的數(shù)據(jù)可捕捉各種物理交互，如剛體動力學（例如運動和碰撞）或光在環(huán)境中的交互。圖 6 顯示了可以從 Replicator 導出的帶注釋數(shù)據(jù)類型示例。

圖 6. Replicator 可以生成各種帶注釋的數(shù)據(jù)，包括(從左上角順時針開始) Normals, RGB, Depth, 和 Semantic Segmentation。

使用 ComfyUI 增強合成數(shù)據(jù)集

進一步增強合成數(shù)據(jù)集可以創(chuàng)建新的變體，例如更改背景以及添加額外的紋理和材質細節(jié)，所有這些操作都使用文本提示。這樣可以生成廣泛的變體，不僅有益，還可以縮小外觀領域差距，獲得高度逼真的結果?？傮w而言，在數(shù)據(jù)集上花費的時間、大型高質量數(shù)字資產(chǎn)的負擔，以及在增強階段重新生成新變體的靈活性可節(jié)省大量時間。

ComfyUI 是一個基于 Web 的后端和 GUI，用于使用擴散模型構建和執(zhí)行管線。這是一個功能強大的開源工具，可在 GitHub 上下載。它可以與 SDXL、其他擴散模型或您選擇的微調模型一起使用。更廣泛的社區(qū)為 ComfyUI 創(chuàng)建了一個包含額外技術和功能的生態(tài)系統(tǒng)，作為 ComfyUI 的節(jié)點，您可以將其集成到您的圖表中。

有關 ComfyUI 圖形的入門參考，請參閱Generative AI for Digital Twins Guide：https://docs.omniverse.nvidia.com/genai-nim-digitaltwin.html

圖 7. 使用擴散模型增強合成生成圖像的 ComfyUI 工作流程示例

從高層次上講，該圖形可被視為“基于模型的編程”。區(qū)域提示用于引導數(shù)據(jù)集圖像中的擴散生成輸出。ControlNet 節(jié)點接收從法線和分割創(chuàng)建的輪廓圖像。ControlNet 與區(qū)域提示相結合，是控制變化的關鍵因素，同時也保留了此數(shù)據(jù)集一致性的重要結構。

啟動‘queue prompt’后，在圖 7 的最右側可以看到增強輸出。這些輸出是傳統(tǒng)渲染的合成數(shù)據(jù)和增強區(qū)域的組合。模型通常能夠從更廣泛的圖像中獲取照明線索和上下文，并對增強區(qū)域進行適當?shù)恼彰骰蜿幱疤幚怼?/p>

地板類型和物體顏色等細節(jié)可以在現(xiàn)有圖像中進行更改。以下四個提示的結果圖像如圖 8 所示。

提示 1

白色平鋪油氈地板
綠色閃亮的新平衡重叉車

木托盤淺色松木、軟木
垃圾容器

提示 2

深色開裂的臟混凝土地板
黃色平衡重叉車
木托盤淺色松木、軟木
和黑色垃圾桶

提示 3

開裂的混凝土地板
白色平衡重叉車
木托盤淺色松木、軟木
垃圾容器

提示 4

綠色切口油氈地板
藍色生銹平衡重叉車
木托盤淺色松木、軟木
垃圾容器

訓練模型

雖然本文并未明確介紹，但下一步邏輯上應該是訓練計算機視覺模型。您可以啟動 NVIDIA 預訓練計算機視覺模型或選擇您自己的模型。然后，可以使用 NVIDIA TAO 等訓練框架對預訓練模型進行微調。TAO 基于 TensorFlow 和 PyTorch 構建，并使用遷移學習來加速模型訓練過程。與使用真實數(shù)據(jù)一樣，您可能必須經(jīng)過多次迭代才能對模型 KPI 感到滿意。

如果您已經(jīng)設置了管線，就可以返回 3D 仿真環(huán)境，通過更改附加參數(shù)和使用 ComfyUI 工作流增強參數(shù)來生成新數(shù)據(jù)。自動化管線減少了生成和標注新數(shù)據(jù)所需的時間，而這通常是訓練模型時的瓶頸。

總結

正如本文所述，您可以使用 NVIDIA NIM 微服務和 NVIDIA Omniverse Replicator 快速輕松地構建自定義合成數(shù)據(jù)生成管線。然后，您可以使用 ComfyUI 進一步增強生成的圖像。這些生成的圖像隨后可以與預訓練模型和工具，如 NVIDIA TAO、PyTorch 或 TensorFlow 等一起使用。

我們很期待看到您如何使用這一工作流程來開發(fā)自己的 SDG 管道。要開始使用，請查看詳細的端到端工作流。