Sensor Simulation

Multi-Camera

The script below is an example of accessing synthetic data from multiple cameras in a scene. It redundantly demonstrates how to access the data either through a custom replicator writer or directly through the built-in annotators.

from omni.isaac.kit import SimulationApp

simulation_app = SimulationApp({"headless": False})

import os

import numpy as np
import omni.kit
import omni.replicator.core as rep
from omni.replicator.core import AnnotatorRegistry, Writer
from PIL import Image
from pxr import Sdf, UsdGeom

NUM_FRAMES = 5

# Save rgb image to file
def save_rgb(rgb_data, file_name):
    rgb_img = Image.fromarray(rgb_data, "RGBA")
    rgb_img.save(file_name + ".png")


# Randomize cube color every frame using a replicator randomizer
def cube_color_randomizer():
    cube_prims = rep.get.prims(path_pattern="Cube")
    with cube_prims:
        rep.randomizer.color(colors=rep.distribution.uniform((0, 0, 0), (1, 1, 1)))
    return cube_prims.node


# Access data through a custom replicator writer
class MyWriter(Writer):
    def __init__(self, rgb: bool = True):
        self._frame_id = 0
        if rgb:
            self.annotators.append(AnnotatorRegistry.get_annotator("rgb"))
        # Create writer output directory
        self.file_path = os.path.join(os.getcwd(), "_out_writer", "")
        print(f"Writing writer data to {self.file_path}")
        dir = os.path.dirname(self.file_path)
        os.makedirs(dir, exist_ok=True)

    def write(self, data):
        for annotator in data.keys():
            annotator_split = annotator.split("-")
            if len(annotator_split) > 1:
                render_product_name = annotator_split[-1]
            if annotator.startswith("rgb"):
                save_rgb(data[annotator], f"{self.file_path}/{render_product_name}_frame_{self._frame_id}")
        self._frame_id += 1


rep.WriterRegistry.register(MyWriter)

# Create a new stage and add a dome light
omni.usd.get_context().new_stage()
stage = omni.usd.get_context().get_stage()
dome_light = stage.DefinePrim("/World/DomeLight", "DomeLight")
dome_light.CreateAttribute("inputs:intensity", Sdf.ValueTypeNames.Float).Set(1000.0)

# Create cube
cube_prim = stage.DefinePrim("/World/Cube", "Cube")
UsdGeom.Xformable(cube_prim).AddTranslateOp().Set((0.0, 5.0, 1.0))

# Run a few frames to ensure all materials are properly loaded
for _ in range(5):
    simulation_app.update()

# Register cube color randomizer to trigger on every frame
rep.randomizer.register(cube_color_randomizer)
with rep.trigger.on_frame():
    rep.randomizer.cube_color_randomizer()

# Create cameras
camera_prim1 = stage.DefinePrim("/World/Camera1", "Camera")
UsdGeom.Xformable(camera_prim1).AddTranslateOp().Set((0.0, 10.0, 20.0))
UsdGeom.Xformable(camera_prim1).AddRotateXYZOp().Set((-15.0, 0.0, 0.0))

camera_prim2 = stage.DefinePrim("/World/Camera2", "Camera")
UsdGeom.Xformable(camera_prim2).AddTranslateOp().Set((-10.0, 15.0, 15.0))
UsdGeom.Xformable(camera_prim2).AddRotateXYZOp().Set((-45.0, 0.0, 45.0))

# Create render products
rp1 = rep.create.render_product(str(camera_prim1.GetPrimPath()), resolution=(320, 320), name="rp1")
rp2 = rep.create.render_product(str(camera_prim2.GetPrimPath()), resolution=(640, 640), name="rp2")
rp3 = rep.create.render_product("/OmniverseKit_Persp", (1024, 1024), name="rp3")

# Acess the data through a custom writer
writer = rep.WriterRegistry.get("MyWriter")
writer.initialize(rgb=True)
writer.attach([rp1, rp2, rp3])

# Acess the data through annotators
rgb_annotators = []
for rp in [rp1, rp2, rp3]:
    rgb = rep.AnnotatorRegistry.get_annotator("rgb")
    rgb.attach(rp)
    rgb_annotators.append(rgb)

# Create annotator output directory
file_path = os.path.join(os.getcwd(), "_out_annot", "")
print(f"Writing annotator data to {file_path}")
dir = os.path.dirname(file_path)
os.makedirs(dir, exist_ok=True)

for i in range(NUM_FRAMES):
    # Improve the rendering quality by rendering multiple subframes per captured frame
    rep.orchestrator.step(rt_subframes=8)
    # Get annotator data after each replicator process step
    for j, rgb_annot in enumerate(rgb_annotators):
        save_rgb(rgb_annot.get_data(), f"{dir}/rp{j}_step_{i}")

simulation_app.close()

Get synthetic data at custom events in simulations

The script below will use annotators and writers (redundantly) to access synthetic data at specific timepoints during a simulation.

from omni.isaac.kit import SimulationApp

simulation_app = SimulationApp(launch_config={"renderer": "RayTracedLighting", "headless": False})

import json
import os

import carb.settings
import numpy as np
import omni
import omni.replicator.core as rep
from omni.isaac.core import World
from omni.isaac.core.objects import DynamicCuboid
from omni.isaac.core.utils.semantics import add_update_semantics
from PIL import Image


# Util function to save rgb annotator data
def write_rgb_data(rgb_data, file_path):
    rgb_img = Image.fromarray(rgb_data, "RGBA")
    rgb_img.save(file_path + ".png")


# Util function to save semantic segmentation annotator data
def write_sem_data(sem_data, file_path):
    id_to_labels = sem_data["info"]["idToLabels"]
    with open(file_path + ".json", "w") as f:
        json.dump(id_to_labels, f)
    sem_image_data = np.frombuffer(sem_data["data"], dtype=np.uint8).reshape(*sem_data["data"].shape, -1)
    sem_img = Image.fromarray(sem_image_data, "RGBA")
    sem_img.save(file_path + ".png")


# Create a new stage with the default ground plane
omni.usd.get_context().new_stage()

# Setup the simulation world
world = World()
world.scene.add_default_ground_plane()
world.reset()

# Setting capture on play to False will prevent the replicator from capturing data each frame
carb.settings.get_settings().set("/omni/replicator/captureOnPlay", False)

# Run a few frames to ensure all materials are properly loaded
for _ in range(5):
    simulation_app.update()

# Create a camera and render product to collect the data from
cam = rep.create.camera(position=(5, 5, 5), look_at=(0, 0, 0))
rp = rep.create.render_product(cam, (512, 512))

# Set the output directory for the data
out_dir = os.getcwd() + "/_out_sim_event"
os.makedirs(out_dir, exist_ok=True)
print(f"Outputting data to {out_dir}..")

# Example of using a writer to save the data
writer = rep.WriterRegistry.get("BasicWriter")
writer.initialize(
    output_dir=f"{out_dir}/writer", rgb=True, semantic_segmentation=True, colorize_semantic_segmentation=True
)
writer.attach(rp)

# Run a preview to ensure the replicator graph is initialized
rep.orchestrator.preview()

# Example of accesing the data directly from annotators
rgb_annot = rep.AnnotatorRegistry.get_annotator("rgb")
rgb_annot.attach(rp)
sem_annot = rep.AnnotatorRegistry.get_annotator("semantic_segmentation", init_params={"colorize": True})
sem_annot.attach(rp)

# Spawn and drop a few cubes, capture data when they stop moving
for i in range(5):
    cuboid = world.scene.add(DynamicCuboid(prim_path=f"/World/Cuboid_{i}", name=f"Cuboid_{i}", position=(0, 0, 10 + i)))
    add_update_semantics(cuboid.prim, "Cuboid")

    for s in range(500):
        world.step(render=False)
        vel = np.linalg.norm(cuboid.get_linear_velocity())
        if vel < 0.1:
            print(f"Cube_{i} stopped moving after {s} simulation steps, writing data..")
            # Tigger the writer and update the annotators with new data
            rep.orchestrator.step(rt_subframes=4, pause_timeline=False)
            write_rgb_data(rgb_annot.get_data(), f"{out_dir}/Cube_{i}_step_{s}_rgb")
            write_sem_data(sem_annot.get_data(), f"{out_dir}/Cube_{i}_step_{s}_sem")
            break

simulation_app.close()

Get Camera Parameters

The below script show how to get the camera parameters associated with a viewport.

import omni
from omni.syntheticdata import helpers
import math

stage = omni.usd.get_context().get_stage()
viewport_api = omni.kit.viewport.utility.get_active_viewport()
# Set viewport resolution, changes will occur on next frame
viewport_api.set_texture_resolution((512, 512))
# get resolution
(width, height) = viewport_api.get_texture_resolution()
aspect_ratio = width / height
# get camera prim attached to viewport
camera = stage.GetPrimAtPath(viewport_api.get_active_camera())
focal_length = camera.GetAttribute("focalLength").Get()
horiz_aperture = camera.GetAttribute("horizontalAperture").Get()
vert_aperture = camera.GetAttribute("verticalAperture").Get()
# Pixels are square so we can also do:
# vert_aperture = height / width * horiz_aperture
near, far = camera.GetAttribute("clippingRange").Get()
fov = 2 * math.atan(horiz_aperture / (2 * focal_length))
# helper to compute projection matrix
proj_mat = helpers.get_projection_matrix(fov, aspect_ratio, near, far)

# compute focal point and center
focal_x = height * focal_length / vert_aperture
focal_y = width * focal_length / horiz_aperture
center_x = height * 0.5
center_y = width * 0.5