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.
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.
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.
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))