diff --git a/icar_sim_sdg/entrypoint.py b/icar_sim_sdg/entrypoint.py index 9e7c2fe..8fbfa12 100644 --- a/icar_sim_sdg/entrypoint.py +++ b/icar_sim_sdg/entrypoint.py @@ -197,10 +197,10 @@ def close(self) -> None: randomizer = base_randomizer.Randomizer() - cam_randomizer = camera_randomizer.CameraRandomizer() - robot_anim_randomizer = robot_randomizer.RobotRandomizer() - scene_props_randomizer = props_randomizer.PropsRandomizer() - scene_lights_randomizer = lights_randomizer.LightsRandomizer() + cam_randomizer = camera_randomizer.CameraRandomizer(randomize_all_frames=not args.cosmos) + robot_anim_randomizer = robot_randomizer.RobotRandomizer(randomize_all_frames=not args.cosmos) + scene_props_randomizer = props_randomizer.PropsRandomizer(randomize_all_frames=not args.cosmos) + scene_lights_randomizer = lights_randomizer.LightsRandomizer(randomize_all_frames=not args.cosmos) randomizer.add_randomizer(cam_randomizer) randomizer.add_randomizer(robot_anim_randomizer) diff --git a/icar_sim_sdg/randomizers/camera_randomizer.py b/icar_sim_sdg/randomizers/camera_randomizer.py index 8afa313..08f29b0 100644 --- a/icar_sim_sdg/randomizers/camera_randomizer.py +++ b/icar_sim_sdg/randomizers/camera_randomizer.py @@ -46,8 +46,17 @@ class CameraRandomizer: """Randomizer for camera position and orientation in the scene.""" - def __init__(self) -> None: - """Initialize the camera randomizer with configuration and camera setup.""" + def __init__(self, randomize_all_frames: bool = True) -> None: + """Initialize the camera randomizer with configuration and camera setup. + + Args: + ---- + randomize_all_frames: If True, randomizes camera pose on every trigger. + If False, randomizes camera pose once during initialization. + + """ + self.randomize_all_frames = randomize_all_frames + self.config = get_config("camera") self.camera, camera_prim_path = create_camera( self.config.get("position", DEFAULT_CAMERA_POSITION), @@ -76,9 +85,13 @@ def __init__(self) -> None: self.altitude_range = self.config.get("camera_altitude", [0.0, 60.0]) # in degrees self.max_sampling_attempts = 5 # Max attempts to find valid position + if not self.randomize_all_frames: + self.update_camera_pose() + def randomize(self) -> None: """Execute camera randomization by updating its pose.""" - self.update_camera_pose() + if self.randomize_all_frames: + self.update_camera_pose() def update_camera_pose( self, diff --git a/icar_sim_sdg/randomizers/lights_randomizer.py b/icar_sim_sdg/randomizers/lights_randomizer.py index 9c0dd2f..91979a9 100644 --- a/icar_sim_sdg/randomizers/lights_randomizer.py +++ b/icar_sim_sdg/randomizers/lights_randomizer.py @@ -28,23 +28,32 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import numpy as np +import omni import omni.replicator.core as rep +from omni.replicator.core.functional.get import light as f_light +from omni.replicator.core.functional.modify import attribute as f_attribute +from pxr import Gf from randomizers.utils import get_config class LightsRandomizer: - """Randomizer for robot animation and pose.""" + """Randomizer for scene lighting attributes.""" - def __init__(self) -> None: - """Initialize the lights randomizer with animation and pose randomization.""" + def __init__(self, randomize_all_frames: bool = True) -> None: + """Initialize the lights randomizer with lighting configuration. + + Args: + ---- + randomize_all_frames: If True, randomizes light attributes on every trigger. + If False, randomizes light attributes once during initialization. + + """ lights_config = get_config("lights") if lights_config.get("enable_lights_randomization", False) is False: return - # Get spawned lights - self.lights = rep.get.light() - # Configuration parameters self.enable_color_temperature = lights_config.get("enable_color_temperature", True) self.min_color = lights_config.get("min_color", [0.0, 0.0, 0.0]) @@ -52,14 +61,31 @@ def __init__(self) -> None: self.min_exposure = lights_config.get("min_exposure", 0.0) self.max_exposure = lights_config.get("max_exposure", 0.5) - def randomize_lights() -> None: - """Randomize lights attributes.""" - with self.lights: - rep.modify.attribute("inputs:color", rep.distribution.uniform(self.min_color, self.max_color)) - rep.modify.attribute("inputs:enableColorTemperature", self.enable_color_temperature) - rep.modify.attribute("inputs:exposure", rep.distribution.uniform(self.min_exposure, self.max_exposure)) + if randomize_all_frames: + # Get spawned lights + self.lights = rep.get.light() + + def randomize_lights() -> None: + """Randomize lights attributes.""" + with self.lights: + rep.modify.attribute("inputs:color", rep.distribution.uniform(self.min_color, self.max_color)) + rep.modify.attribute("inputs:enableColorTemperature", self.enable_color_temperature) + rep.modify.attribute( + "inputs:exposure", rep.distribution.uniform(self.min_exposure, self.max_exposure) + ) - rep.randomizer.register(randomize_lights) + rep.randomizer.register(randomize_lights) - with rep.trigger.on_frame(): - randomize_lights() + with rep.trigger.on_frame(): + randomize_lights() + else: + stage = omni.usd.get_context().get_stage() + lights = f_light(stage=stage) + for light in lights: + f_attribute( + light, "inputs:color", Gf.Vec3f(*np.random.uniform(low=self.min_color, high=self.max_color)) + ) + f_attribute(light, "inputs:enableColorTemperature", self.enable_color_temperature) + f_attribute( + light, "inputs:exposure", float(np.random.uniform(low=self.min_exposure, high=self.max_exposure)) + ) diff --git a/icar_sim_sdg/randomizers/props_randomizer.py b/icar_sim_sdg/randomizers/props_randomizer.py index 6a9cda8..c985f6f 100644 --- a/icar_sim_sdg/randomizers/props_randomizer.py +++ b/icar_sim_sdg/randomizers/props_randomizer.py @@ -30,7 +30,11 @@ import pathlib +import numpy as np import omni.replicator.core as rep +from omni.replicator.core.functional.create import plane as f_plane +from omni.replicator.core.functional.create import reference as f_reference +from omni.replicator.core.functional.randomizer import scatter_2d as f_scatter_2d from randomizers.utils import get_config PROPS_PATH = pathlib.Path(__file__).parent.parent.parent / "assets" / "props" @@ -39,9 +43,16 @@ class PropsRandomizer: """Randomizer for scattering props on a surface in the scene.""" - def __init__(self) -> None: - """Initialize the props randomizer with configuration.""" - self.surface = rep.create.plane(position=(0, 0, 0), scale=0.75, visible=False) + def __init__(self, randomize_all_frames: bool = True) -> None: + """Initialize the props randomizer with configuration. + + Args: + ---- + randomize_all_frames: If True, randomizes prop positions and rotations on every trigger. + If False, randomizes prop positions and rotations once during initialization. + + """ + self.randomize_all_frames = randomize_all_frames self.config = get_config("props") self.num_props = self.config.get("num_props", 10) @@ -70,15 +81,23 @@ def __init__(self) -> None: self.rotation_min, self.rotation_max = (self.rotation_range["z"][0],), (self.rotation_range["z"][1],) self.usds = rep.utils.get_usd_files(str(PROPS_PATH), recursive=True) - def randomize_props() -> None: - """Randomize prop positions and rotations.""" - with rep.randomizer.instantiate(self.usds, size=self.num_props, semantics=[("class", "prop")]): - rep.randomizer.scatter_2d( - check_for_collisions=True, - surface_prims=self.surface, - ) + if self.randomize_all_frames: + surface = rep.create.plane(position=(0, 0, 0), scale=0.75, visible=False) + + def randomize_props() -> None: + """Randomize prop positions and rotations.""" + with rep.randomizer.instantiate(self.usds, size=self.num_props, semantics=[("class", "prop")]): + rep.randomizer.scatter_2d( + check_for_collisions=True, + surface_prims=surface, + ) - rep.randomizer.register(randomize_props) + rep.randomizer.register(randomize_props) - with rep.trigger.on_frame(): - randomize_props() + with rep.trigger.on_frame(): + randomize_props() + else: + surface = f_plane(position=(0, 0, 0), scale=0.75, visible=False) + random_idx = np.random.randint(0, len(self.usds) - 1, size=self.num_props) + props_prims = [f_reference(usd_path=self.usds[idx], semantics={"class": "prop"}) for idx in random_idx] + f_scatter_2d(prims=props_prims, surface_prims=surface, check_for_collisions=True) diff --git a/icar_sim_sdg/randomizers/robot_randomizer.py b/icar_sim_sdg/randomizers/robot_randomizer.py index 7072aac..bcb997a 100644 --- a/icar_sim_sdg/randomizers/robot_randomizer.py +++ b/icar_sim_sdg/randomizers/robot_randomizer.py @@ -30,27 +30,61 @@ import pathlib +import numpy as np import omni.replicator.core as rep +import omni.timeline +from omni.replicator.core.functional.create import reference as f_reference ROBOT_PATH = pathlib.Path(__file__).parent.parent.parent / "assets" / "ur10e" / "pick_animation" / "playback.usda" +FIRST_MOVING_FRAME = 80 NUM_FRAMES = 642 # Animation frames class RobotRandomizer: """Randomizer for robot animation and pose.""" - def __init__(self) -> None: - """Initialize the robot randomizer with animation and pose randomization.""" - # Spawn robot pick animation - self.robot = rep.create.from_usd(str(ROBOT_PATH), semantics={"class": "robot"}) + def __init__(self, randomize_all_frames: bool = True) -> None: + """Initialize the robot randomizer with animation and pose randomization. - def randomize_robot_pose() -> None: - """Randomize robot animation frame and rotation.""" - with self.robot: - rep.modify.timeline(rep.distribution.uniform(lower=0, upper=NUM_FRAMES), modify_type="frame") - rep.modify.pose(rotation_z=rep.distribution.uniform(0.0, 360.0)) + Args: + ---- + randomize_all_frames: If True, randomizes animation frame and rotation on every trigger. + If False, incrementally advances through animation frames. - rep.randomizer.register(randomize_robot_pose) + """ + self.randomize_all_frames = randomize_all_frames + self.timeline_iface = omni.timeline.acquire_timeline_interface() + self.current_frame = FIRST_MOVING_FRAME - with rep.trigger.on_frame(): - randomize_robot_pose() + if self.randomize_all_frames: + # Spawn robot pick animation + self.robot = rep.create.from_usd(str(ROBOT_PATH), semantics={"class": "robot"}) + + def randomize_robot_pose() -> None: + """Randomize robot animation frame and rotation.""" + with self.robot: + rep.modify.timeline(rep.distribution.uniform(lower=0, upper=NUM_FRAMES), modify_type="frame") + rep.modify.pose(rotation_z=rep.distribution.uniform(0.0, 360.0)) + + rep.randomizer.register(randomize_robot_pose) + + with rep.trigger.on_frame(): + randomize_robot_pose() + else: + random_rotation = np.random.uniform(0.0, 360.0) + f_reference( + usd_path=str(ROBOT_PATH), + position=(0, 0, 0), + rotation=(0, 0, random_rotation), + semantics={"class": "robot"}, + ) + + def randomize(self) -> None: + """Apply randomization to the robot.""" + if self.randomize_all_frames: + return + fps = float(self.timeline_iface.get_time_codes_per_seconds()) + self.timeline_iface.set_current_time(self.current_frame / fps) + self.current_frame += 1 + if self.current_frame >= NUM_FRAMES: + self.current_frame = FIRST_MOVING_FRAME