10.14. Randomization Snippets

Examples of randomization using USD and Isaac Sim APIs. These examples demonstrate how to randomize scenes for synthetic data generation (SDG) in scenarios where default replicator randomizers are not sufficient or applicable.

The snippets are designed to align with the structure and function names used in the replicator example snippets. In comparison they also have the option to write the data to disk by stetting write_data=True.

Prerequisites:

  • Familiarity with USD.

  • Ability to execute code from the Script Editor.

  • Understanding basic replicator concepts, such as subframes.

10.14.1. Randomizing a light source

This snippet sets up a new environment containing a cube and a sphere; it then spawns a given number of lights and randomizes selected attributes for these lights over a specified number of frames.

../_images/isaac_tutorial_replicator_randomization_lights.gif
Randomizing a light source
 1import asyncio
 2import os
 3
 4import numpy as np
 5import omni.kit.commands
 6import omni.replicator.core as rep
 7import omni.usd
 8from omni.isaac.core.utils.semantics import add_update_semantics
 9from pxr import Gf, Sdf, UsdGeom
10
11omni.usd.get_context().new_stage()
12stage = omni.usd.get_context().get_stage()
13
14sphere = stage.DefinePrim("/World/Sphere", "Sphere")
15UsdGeom.Xformable(sphere).AddTranslateOp().Set((0.0, 1.0, 1.0))
16add_update_semantics(sphere, "sphere", "class")
17
18cube = stage.DefinePrim("/World/Cube", "Cube")
19UsdGeom.Xformable(cube).AddTranslateOp().Set((0.0, -2.0, 2.0))
20add_update_semantics(cube, "cube", "class")
21
22plane_path = "/World/Plane"
23omni.kit.commands.execute("CreateMeshPrimWithDefaultXform", prim_path=plane_path, prim_type="Plane")
24plane_prim = stage.GetPrimAtPath(plane_path)
25plane_prim.CreateAttribute("xformOp:scale", Sdf.ValueTypeNames.Double3, False).Set(Gf.Vec3d(10, 10, 1))
26
27
28def sphere_lights(num):
29    lights = []
30    for i in range(num):
31        # "CylinderLight", "DiskLight", "DistantLight", "DomeLight", "RectLight", "SphereLight"
32        prim_type = "SphereLight"
33        next_free_path = omni.usd.get_stage_next_free_path(stage, f"/World/{prim_type}", False)
34        light_prim = stage.DefinePrim(next_free_path, prim_type)
35        UsdGeom.Xformable(light_prim).AddTranslateOp().Set((0.0, 0.0, 0.0))
36        UsdGeom.Xformable(light_prim).AddRotateXYZOp().Set((0.0, 0.0, 0.0))
37        UsdGeom.Xformable(light_prim).AddScaleOp().Set((1.0, 1.0, 1.0))
38        light_prim.CreateAttribute("inputs:enableColorTemperature", Sdf.ValueTypeNames.Bool).Set(True)
39        light_prim.CreateAttribute("inputs:colorTemperature", Sdf.ValueTypeNames.Float).Set(6500.0)
40        light_prim.CreateAttribute("inputs:radius", Sdf.ValueTypeNames.Float).Set(0.5)
41        light_prim.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(30000.0)
42        light_prim.CreateAttribute("inputs:color", Sdf.ValueTypeNames.Color3f).Set((1.0, 1.0, 1.0))
43        light_prim.CreateAttribute("inputs:exposure", Sdf.ValueTypeNames.Float).Set(0.0)
44        light_prim.CreateAttribute("inputs:diffuse", Sdf.ValueTypeNames.Float).Set(1.0)
45        light_prim.CreateAttribute("inputs:specular", Sdf.ValueTypeNames.Float).Set(1.0)
46        lights.append(light_prim)
47    return lights
48
49
50async def run_randomizations_async(num_frames, lights, write_data=True, delay=0):
51    if write_data:
52        writer = rep.WriterRegistry.get("BasicWriter")
53        out_dir = os.getcwd() + "/_out_rand_lights"
54        print(f"Writing data to {out_dir}..")
55        writer.initialize(output_dir=out_dir, rgb=True)
56        rp = rep.create.render_product("/OmniverseKit_Persp", (512, 512))
57        writer.attach(rp)
58
59    for _ in range(num_frames):
60        for light in lights:
61            light.GetAttribute("xformOp:translate").Set(
62                (np.random.uniform(-5, 5), np.random.uniform(-5, 5), np.random.uniform(4, 6))
63            )
64            scale_rand = np.random.uniform(0.5, 1.5)
65            light.GetAttribute("xformOp:scale").Set((scale_rand, scale_rand, scale_rand))
66            light.GetAttribute("inputs:colorTemperature").Set(np.random.normal(4500, 1500))
67            light.GetAttribute("inputs:intensity").Set(np.random.normal(25000, 5000))
68            light.GetAttribute("inputs:color").Set(
69                (np.random.uniform(0.1, 0.9), np.random.uniform(0.1, 0.9), np.random.uniform(0.1, 0.9))
70            )
71
72        if write_data:
73            await rep.orchestrator.step_async(rt_subframes=16)
74        else:
75            await omni.kit.app.get_app().next_update_async()
76        if delay > 0:
77            await asyncio.sleep(delay)
78
79
80num_frames = 10
81lights = sphere_lights(10)
82asyncio.ensure_future(run_randomizations_async(num_frames=num_frames, lights=lights, delay=0.2))

10.14.2. Randomizing textures

The snippet sets up an environment, spawns a given number of cubes and spheres, and randomizes their textures for the given number of frames. After the randomizations their original materials are reassigned. The snippet also showcases how to create a new material and assign it to a prim.

../_images/isaac_tutorial_replicator_randomization_textures.gif
Randomizing textures
  1import asyncio
  2import os
  3
  4import numpy as np
  5import omni.replicator.core as rep
  6import omni.usd
  7from omni.isaac.core.utils.nucleus import get_assets_root_path
  8from omni.isaac.core.utils.semantics import add_update_semantics, get_semantics
  9from pxr import Gf, Sdf, UsdGeom, UsdShade
 10
 11omni.usd.get_context().new_stage()
 12stage = omni.usd.get_context().get_stage()
 13dome_light = stage.DefinePrim("/World/DomeLight", "DomeLight")
 14dome_light.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(1000.0)
 15
 16sphere = stage.DefinePrim("/World/Sphere", "Sphere")
 17UsdGeom.Xformable(sphere).AddTranslateOp().Set((0.0, 0.0, 1.0))
 18add_update_semantics(sphere, "sphere", "class")
 19
 20num_cubes = 10
 21for _ in range(num_cubes):
 22    prim_type = "Cube"
 23    next_free_path = omni.usd.get_stage_next_free_path(stage, f"/World/{prim_type}", False)
 24    cube = stage.DefinePrim(next_free_path, prim_type)
 25    UsdGeom.Xformable(cube).AddTranslateOp().Set((np.random.uniform(-3.5, 3.5), np.random.uniform(-3.5, 3.5), 1))
 26    scale_rand = np.random.uniform(0.25, 0.5)
 27    UsdGeom.Xformable(cube).AddScaleOp().Set((scale_rand, scale_rand, scale_rand))
 28    add_update_semantics(cube, "cube", "class")
 29
 30plane_path = "/World/Plane"
 31omni.kit.commands.execute("CreateMeshPrimWithDefaultXform", prim_path=plane_path, prim_type="Plane")
 32plane_prim = stage.GetPrimAtPath(plane_path)
 33plane_prim.CreateAttribute("xformOp:scale", Sdf.ValueTypeNames.Double3, False).Set(Gf.Vec3d(10, 10, 1))
 34
 35
 36def get_shapes():
 37    stage = omni.usd.get_context().get_stage()
 38    shapes = []
 39    for prim in stage.Traverse():
 40        sem_dict = get_semantics(prim)
 41        sem_values = sem_dict.values()
 42        if ("class", "cube") in sem_values or ("class", "sphere") in sem_values:
 43            shapes.append(prim)
 44    return shapes
 45
 46
 47shapes = get_shapes()
 48
 49
 50def create_omnipbr_material(mtl_url, mtl_name, mtl_path):
 51    stage = omni.usd.get_context().get_stage()
 52    omni.kit.commands.execute("CreateMdlMaterialPrim", mtl_url=mtl_url, mtl_name=mtl_name, mtl_path=mtl_path)
 53    material_prim = stage.GetPrimAtPath(mtl_path)
 54    shader = UsdShade.Shader(omni.usd.get_shader_from_material(material_prim, get_prim=True))
 55
 56    # Add value inputs
 57    shader.CreateInput("diffuse_color_constant", Sdf.ValueTypeNames.Color3f)
 58    shader.CreateInput("reflection_roughness_constant", Sdf.ValueTypeNames.Float)
 59    shader.CreateInput("metallic_constant", Sdf.ValueTypeNames.Float)
 60
 61    # Add texture inputs
 62    shader.CreateInput("diffuse_texture", Sdf.ValueTypeNames.Asset)
 63    shader.CreateInput("reflectionroughness_texture", Sdf.ValueTypeNames.Asset)
 64    shader.CreateInput("metallic_texture", Sdf.ValueTypeNames.Asset)
 65
 66    # Add other attributes
 67    shader.CreateInput("project_uvw", Sdf.ValueTypeNames.Bool)
 68
 69    # Add texture scale and rotate
 70    shader.CreateInput("texture_scale", Sdf.ValueTypeNames.Float2)
 71    shader.CreateInput("texture_rotate", Sdf.ValueTypeNames.Float)
 72
 73    material = UsdShade.Material(material_prim)
 74    return material
 75
 76
 77def create_materials(num):
 78    MDL = "OmniPBR.mdl"
 79    mtl_name, _ = os.path.splitext(MDL)
 80    MAT_PATH = "/World/Looks"
 81    materials = []
 82    for _ in range(num):
 83        prim_path = omni.usd.get_stage_next_free_path(stage, f"{MAT_PATH}/{mtl_name}", False)
 84        mat = create_omnipbr_material(mtl_url=MDL, mtl_name=mtl_name, mtl_path=prim_path)
 85        materials.append(mat)
 86    return materials
 87
 88
 89materials = create_materials(len(shapes))
 90
 91
 92async def run_randomizations_async(num_frames, materials, textures, write_data=True, delay=0):
 93    if write_data:
 94        writer = rep.WriterRegistry.get("BasicWriter")
 95        out_dir = os.getcwd() + "/_out_rand_textures"
 96        print(f"Writing data to {out_dir}..")
 97        writer.initialize(output_dir=out_dir, rgb=True)
 98        rp = rep.create.render_product("/OmniverseKit_Persp", (512, 512))
 99        writer.attach(rp)
100
101    # Apply the new materials and store the initial ones to reassign later
102    initial_materials = {}
103    for i, shape in enumerate(shapes):
104        cur_mat, _ = UsdShade.MaterialBindingAPI(shape).ComputeBoundMaterial()
105        initial_materials[shape] = cur_mat
106        UsdShade.MaterialBindingAPI(shape).Bind(materials[i], UsdShade.Tokens.strongerThanDescendants)
107
108    for _ in range(num_frames):
109        for mat in materials:
110            shader = UsdShade.Shader(omni.usd.get_shader_from_material(mat, get_prim=True))
111            diffuse_texture = np.random.choice(textures)
112            shader.GetInput("diffuse_texture").Set(diffuse_texture)
113            project_uvw = np.random.choice([True, False], p=[0.9, 0.1])
114            shader.GetInput("project_uvw").Set(bool(project_uvw))
115            texture_scale = np.random.uniform(0.1, 1)
116            shader.GetInput("texture_scale").Set((texture_scale, texture_scale))
117            texture_rotate = np.random.uniform(0, 45)
118            shader.GetInput("texture_rotate").Set(texture_rotate)
119
120        if write_data:
121            await rep.orchestrator.step_async(rt_subframes=4)
122        else:
123            await omni.kit.app.get_app().next_update_async()
124        if delay > 0:
125            await asyncio.sleep(delay)
126
127    # Reassign the initial materials
128    for shape, mat in initial_materials.items():
129        if mat:
130            UsdShade.MaterialBindingAPI(shape).Bind(mat, UsdShade.Tokens.strongerThanDescendants)
131        else:
132            UsdShade.MaterialBindingAPI(shape).UnbindAllBindings()
133
134
135assets_root_path = get_assets_root_path()
136textures = [
137    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/aggregate_exposed_diff.jpg",
138    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/gravel_track_ballast_diff.jpg",
139    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/gravel_track_ballast_multi_R_rough_G_ao.jpg",
140    assets_root_path + "/NVIDIA/Materials/vMaterials_2/Ground/textures/rough_gravel_rough.jpg",
141]
142
143num_frames = 10
144asyncio.ensure_future(run_randomizations_async(num_frames, materials, textures, delay=0.2))

10.14.3. Chained randomizations

The snippet gives an example of a more complex randomizations, where the results of the first randomization are used to determine the next randomization. It uses a custom sampler function to set the location of the camera by iterating (almost) equidistant points on a sphere. The snippet starts by setting up the environment, a forklift, a pallet, a bin, and a dome light. For every randomization frame it will cycle through the dome light textures, move the pallet in a random location, on top of which it will move the bin so that it is fully on top of the pallet. Finally it will move the camera to a new location on a sphere looking at the bin.

../_images/isaac_tutorial_replicator_randomization_chained_persp.gif ../_images/isaac_tutorial_replicator_randomization_chained_sphere.gif
Chained randomizations
  1import asyncio
  2import itertools
  3import os
  4
  5import numpy as np
  6import omni.replicator.core as rep
  7import omni.usd
  8from omni.isaac.core.utils.nucleus import get_assets_root_path
  9from pxr import Gf, Usd, UsdGeom, UsdLux
 10
 11
 12# https://stackoverflow.com/questions/9600801/evenly-distributing-n-points-on-a-sphere
 13# https://arxiv.org/pdf/0912.4540.pdf
 14def next_point_on_sphere(idx, num_points, radius=1, origin=(0, 0, 0)):
 15    offset = 2.0 / num_points
 16    inc = np.pi * (3.0 - np.sqrt(5.0))
 17    z = ((idx * offset) - 1) + (offset / 2)
 18    phi = ((idx + 1) % num_points) * inc
 19    r = np.sqrt(1 - pow(z, 2))
 20    y = np.cos(phi) * r
 21    x = np.sin(phi) * r
 22    return [(x * radius) + origin[0], (y * radius) + origin[1], (z * radius) + origin[2]]
 23
 24
 25assets_root_path = get_assets_root_path()
 26FORKLIFT_PATH = assets_root_path + "/Isaac/Props/Forklift/forklift.usd"
 27PALLET_PATH = assets_root_path + "/Isaac/Props/Pallet/pallet.usd"
 28BIN_PATH = assets_root_path + "/Isaac/Props/KLT_Bin/small_KLT_visual.usd"
 29
 30omni.usd.get_context().new_stage()
 31stage = omni.usd.get_context().get_stage()
 32
 33dome_light = UsdLux.DomeLight.Define(stage, "/World/Lights/DomeLight")
 34dome_light.GetIntensityAttr().Set(1000)
 35
 36forklift_prim = stage.DefinePrim("/World/Forklift", "Xform")
 37forklift_prim.GetReferences().AddReference(FORKLIFT_PATH)
 38if not forklift_prim.GetAttribute("xformOp:translate"):
 39    UsdGeom.Xformable(forklift_prim).AddTranslateOp()
 40forklift_prim.GetAttribute("xformOp:translate").Set((-4.5, -4.5, 0))
 41
 42pallet_prim = stage.DefinePrim("/World/Pallet", "Xform")
 43pallet_prim.GetReferences().AddReference(PALLET_PATH)
 44if not pallet_prim.GetAttribute("xformOp:translate"):
 45    UsdGeom.Xformable(pallet_prim).AddTranslateOp()
 46if not pallet_prim.GetAttribute("xformOp:rotateXYZ"):
 47    UsdGeom.Xformable(pallet_prim).AddRotateXYZOp()
 48
 49bin_prim = stage.DefinePrim("/World/Bin", "Xform")
 50bin_prim.GetReferences().AddReference(BIN_PATH)
 51if not bin_prim.GetAttribute("xformOp:translate"):
 52    UsdGeom.Xformable(bin_prim).AddTranslateOp()
 53if not bin_prim.GetAttribute("xformOp:rotateXYZ"):
 54    UsdGeom.Xformable(bin_prim).AddRotateXYZOp()
 55
 56cam = stage.DefinePrim("/World/Camera", "Camera")
 57if not cam.GetAttribute("xformOp:translate"):
 58    UsdGeom.Xformable(cam).AddTranslateOp()
 59if not cam.GetAttribute("xformOp:orient"):
 60    UsdGeom.Xformable(cam).AddOrientOp()
 61
 62
 63async def run_randomizations_async(
 64    num_frames, dome_light, dome_textures, pallet_prim, bin_prim, write_data=True, delay=0
 65):
 66    if write_data:
 67        writer = rep.WriterRegistry.get("BasicWriter")
 68        out_dir = os.getcwd() + "/_out_rand_sphere_scan"
 69        print(f"Writing data to {out_dir}..")
 70        writer.initialize(output_dir=out_dir, rgb=True)
 71        rp_persp = rep.create.render_product("/OmniverseKit_Persp", (512, 512), name="PerspView")
 72        rp_cam = rep.create.render_product(str(cam.GetPath()), (512, 512), name="SphereView")
 73        writer.attach([rp_cam, rp_persp])
 74
 75    textures_cycle = itertools.cycle(dome_textures)
 76
 77    bb_cache = UsdGeom.BBoxCache(time=Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_])
 78    pallet_size = bb_cache.ComputeWorldBound(pallet_prim).GetRange().GetSize()
 79    pallet_length = pallet_size.GetLength()
 80    bin_size = bb_cache.ComputeWorldBound(bin_prim).GetRange().GetSize()
 81
 82    for i in range(num_frames):
 83        # Set next background texture every nth frame and run an app update
 84        if i % 5 == 0:
 85            dome_light.GetTextureFileAttr().Set(next(textures_cycle))
 86            await omni.kit.app.get_app().next_update_async()
 87
 88        # Randomize pallet pose
 89        pallet_prim.GetAttribute("xformOp:translate").Set(
 90            Gf.Vec3d(np.random.uniform(-1.5, 1.5), np.random.uniform(-1.5, 1.5), 0)
 91        )
 92        rand_z_rot = np.random.uniform(-90, 90)
 93        pallet_prim.GetAttribute("xformOp:rotateXYZ").Set(Gf.Vec3d(0, 0, rand_z_rot))
 94        pallet_tf_mat = omni.usd.get_world_transform_matrix(pallet_prim)
 95        pallet_rot = pallet_tf_mat.ExtractRotation()
 96        pallet_pos = pallet_tf_mat.ExtractTranslation()
 97
 98        # Randomize bin position on top of the rotated pallet area making sure the bin is fully on the pallet
 99        rand_transl_x = np.random.uniform(-pallet_size[0] / 2 + bin_size[0] / 2, pallet_size[0] / 2 - bin_size[0] / 2)
100        rand_transl_y = np.random.uniform(-pallet_size[1] / 2 + bin_size[1] / 2, pallet_size[1] / 2 - bin_size[1] / 2)
101
102        # Adjust bin position to account for the random rotation of the pallet
103        rand_z_rot_rad = np.deg2rad(rand_z_rot)
104        rot_adjusted_transl_x = rand_transl_x * np.cos(rand_z_rot_rad) - rand_transl_y * np.sin(rand_z_rot_rad)
105        rot_adjusted_transl_y = rand_transl_x * np.sin(rand_z_rot_rad) + rand_transl_y * np.cos(rand_z_rot_rad)
106        bin_prim.GetAttribute("xformOp:translate").Set(
107            Gf.Vec3d(
108                pallet_pos[0] + rot_adjusted_transl_x,
109                pallet_pos[1] + rot_adjusted_transl_y,
110                pallet_pos[2] + pallet_size[2] + bin_size[2] / 2,
111            )
112        )
113        # Keep bin rotation aligned with pallet
114        bin_prim.GetAttribute("xformOp:rotateXYZ").Set(pallet_rot.GetAxis() * pallet_rot.GetAngle())
115
116        # Get next camera position on a sphere looking at the bin with a randomized distance
117        rand_radius = np.random.normal(3, 0.5) * pallet_length
118        bin_pos = omni.usd.get_world_transform_matrix(bin_prim).ExtractTranslation()
119        cam_pos = next_point_on_sphere(i, num_points=num_frames, radius=rand_radius, origin=bin_pos)
120        cam.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*cam_pos))
121
122        eye = Gf.Vec3d(*cam_pos)
123        target = Gf.Vec3d(*bin_pos)
124        up_axis = Gf.Vec3d(0, 0, 1)
125        look_at_quatd = Gf.Matrix4d().SetLookAt(eye, target, up_axis).GetInverse().ExtractRotation().GetQuat()
126        cam.GetAttribute("xformOp:orient").Set(Gf.Quatf(look_at_quatd))
127
128        if write_data:
129            await rep.orchestrator.step_async(rt_subframes=4)
130        else:
131            await omni.kit.app.get_app().next_update_async()
132        if delay > 0:
133            await asyncio.sleep(delay)
134
135
136num_frames = 90
137dome_textures = [
138    assets_root_path + "/NVIDIA/Assets/Skies/Cloudy/champagne_castle_1_4k.hdr",
139    assets_root_path + "/NVIDIA/Assets/Skies/Clear/evening_road_01_4k.hdr",
140    assets_root_path + "/NVIDIA/Assets/Skies/Clear/mealie_road_4k.hdr",
141    assets_root_path + "/NVIDIA/Assets/Skies/Clear/qwantani_4k.hdr",
142]
143asyncio.ensure_future(run_randomizations_async(num_frames, dome_light, dome_textures, pallet_prim, bin_prim, delay=0.2))