Replicator - nested loops or task queque

[simon.steinmann]

@jlafleche ,

Your response in this post How to get current frame number in Replicator - #7 by jlafleche
helped me A LOT, and I managed to get things running. However, now I am faced with another issue, which probably can be solved in a more elegant manner. Basically I have a scripted camera movement, which moves in a circle around the target and generates a labeled dataset. Now I want to change the angle of the camera, or the distance, or the material target, once a circle is complete. Essentially something like this:

elevation_list = [12, 30, 60, 89.9],
distance_list = [1000, 1600]
color_list = [[0.2806969, 0.2806969, 0.2806969], [0.028681736, 0.104105994, 0.28571427]]

with rep.trigger.on_frame(num_frames=10000):
    for color in color_list:
        for elevation in elevation_list:
            for distance in distance_list:
                with rep.get.prims("/Root/Volvo_S90_Recharge_Free/Materials/Car_Paint/pbr_shader"):
                    rep.modify.attribute("inputs:diffuseColor", color)
                with camera:
                    rep.modify.pose(
                        position=rep.randomizer.compute_camera_pos(elevation, distance),
                        look_at=(0, 0, 0)
                    )

I know this does not work, and I got it working with some very ugly global variables and frame count calculations, but that approach is not very scalable. Something akin a nested for loop like in my code snippet, or some sort of task queque would be great. Basically we want to define our dataset generation in a config file, and just hit “go” and have the entire dataset be generated.

I would really appreciate any pointers or tips.

Thank you,
Simon

[pcallender]

Hi @simon.steinmann , please check out the code snippet below. This works in the current release, but a different solution will be needed for the upcoming release due to changes in the way autofunc works. I’ll follow up with more info after the new release. For now, hopefully, this should unblock you, and be an example for this type of nesting.

import omni.replicator.core as rep
import omni.graph.core as og
import numpy as np
from typing import Any
import carb


rep.settings.set_render_rtx_realtime()

# PROTOTYPE SOLUTION
# Add functionality to ReplicatorItem
def set_input(self, attribute: str, value: Any):
    if self.node is None:
        raise ValueError(f"ReplicatorItem object has no associated node, unable to retrieve attribute {attribute}.")
    if not self.node.get_attribute_exists(f"inputs:{attribute}"):
        raise ValueError(f"ReplicatorItem node has no attribute matching: {attribute}.")
    node_attribute = self.node.get_attribute(f"inputs:{attribute}")
    resolved_type = node_attribute.get_resolved_type()
    # Disconnect any existing connection
    for conn in node_attribute.get_upstream_connections():
        conn.disconnect(node_attribute, False)
    # Set data type
    node_attribute.set_resolved_type(resolved_type)
    # Set attribute
    return node_attribute.set(value)

rep.utils.ReplicatorItem.set_input = set_input

frame = 0

elevation_list = [12, 30, 60, 89.9]
distance_list = [1000, 1600]
color_list = [[0.2806969, 0.2806969, 0.2806969], [0.028681736, 0.104105994, 0.28571427]]

@og.AutoFunc(module_name="omni.replicator")
def ComputeCameraPos(distance: float, elevation: float, azimuth: float, numSamples: int = 1) -> og.Bundle:
    # Note 1: numSamples input currently required
    # Note 2: Only bundle output currently supported, this will be expanded in the future.

    # Use global to have access to a persistent `frame` variable
	global frame
	frame += 1
	x = np.cos(azimuth) * np.cos(elevation) * distance
	y = np.sin(azimuth) * np.cos(elevation) * distance
	z = np.sin(elevation) * distance

	bundle = og.Bundle("return", False)
	bundle.create_attribute("values", og.Type(og.BaseDataType.DOUBLE, 3, 1)).value = [[x, y, z]]
	return bundle

# This will allow the AutoFunc return attribute `out_0` to be automatically connected to the pose node's `values` input
rep.utils.ATTRIBUTE_MAPPINGS.add(rep.utils.AttrMap("outputs_out_0", "inputs:values"))

# Register functionality into replicator
def compute_camera_pos(distance: float, elevation: float, azimuth: float):
	return rep.utils.create_node("omni.replicator.ComputeCameraPos", distance=distance, elevation=elevation, azimuth=azimuth)
rep.randomizer.register(compute_camera_pos)
camera = rep.create.camera()
rp = rep.create.render_product(camera, (1024, 1024))
writer = rep.WriterRegistry.get("BasicWriter")
writer.initialize(output_dir="_forloop", rgb=True)
writer.attach(rp)

rep.create.cube()

with rep.trigger.on_frame(num_frames=10000):
	with rep.get.prims("/Replicator/Cube"):
		color_sampler = rep.distribution.uniform((0., 0., 0.), (1., 1., 1.))
		rep.randomizer.color(color_sampler)
	with camera:
		camera_pose = rep.randomizer.compute_camera_pos(0, 0, 0)
		rep.modify.pose(
		position=camera_pose,
		look_at=(0, 0, 0)
	)

async def main():
    for color in color_list:
        for elevation in elevation_list:
            for distance in distance_list:
                for azimuth in range(0, 10):
                    color_sampler.set_input("lower", color)
                    color_sampler.set_input("upper", color)
                    camera_pose.set_input("elevation", elevation)
                    camera_pose.set_input("distance", distance)
                    camera_pose.set_input("azimuth", azimuth * 3.6)
                    
                    # Render one frame
                    await omni.kit.app.get_app().next_update_async()
                    # Start capture if not yet started
                    if not rep.orchestrator.get_is_started():
                        await rep.orchestrator._Orchestrator()._start_async()
                    
    rep.orchestrator.stop()

import asyncio
asyncio.ensure_future(main())

[simon.steinmann]

Oh wow, this looks very promising. I will certainly check it out. Thank you so much! I hope in the future some default method will be implemented to allow for this type of more deterministic workflow. I’m looking forward to the updates, and if there is anything you want to have beta-tested or need any other help or seek a feature request discussion, just let me know. I am very keen to help!

[simon.steinmann]

First of all, your code works and it makes sense, however I have trouble porting it to other node types. I think I need your help again. I have tried many things, but I cant get this approach working with spawning models and materials. I want something like this

async def main():
    for model in model_list:
        for pose in model["pose_list"]:
            for material in model["material_list"]:
                for color in material["color_list"]:
                    for elevation in elevation_list:
                        for distance in distance_list:
                            for azimuth in range(frames_per_circle):

The color and pose I can figure out, but how to assign a material to a model I can’t figure out. Perhaps setting all models to invisible, and then editing the raw usd properties of the active model. I’m a bit at a loss here…

Another Idea would be using AutoFunc to return a string or prim. But I do not know how to do that via packing it into a bundle, and I could not figure it out with the documentation. If that works, I could perhaps do something like this:

with rep.trigger.on_frame(num_frames=10000):
    with rep.get.prims(semantics=[('class', 'car')]):
        rep.modify.visibility(rep.randomizer.get_visibility_bool_list())
    with rep.randomizer.get_model():        
        rep.randomizer.materials(rep.randomizer.get_material())
        rep.randomizer.color(rep.randomizer.get_color())
        rep.modify.pose(rep.randomizer.get_pose())