47 lines
2.1 KiB
Python
47 lines
2.1 KiB
Python
import os
|
|
|
|
from nimbus.components.data.iterator import Iterator
|
|
from nimbus.components.data.scene import Scene
|
|
from nimbus.components.load import LayoutRandomizer
|
|
|
|
|
|
class EnvRandomizer(LayoutRandomizer):
|
|
"""
|
|
Environment randomizer that extends the base layout randomizer to include additional randomization
|
|
capabilities specific to the simulation environment.
|
|
This class can be used to randomize various aspects of the environment, such as object placements,
|
|
textures, lighting conditions, and other scene parameters, based on the provided configuration.
|
|
The randomization process can be controlled through the number of randomizations to perform and
|
|
whether to operate in strict mode.
|
|
|
|
Args:
|
|
scene_iter (Iterator[Scene]): An iterator that yields scenes to be randomized.
|
|
random_num (int): How many randomizations to perform for each scene.
|
|
strict_mode (bool): Whether to operate in strict mode, which enforces certain constraints
|
|
on the randomization process.
|
|
input_dir (str): Directory from which to load additional randomization data such as object
|
|
placements or textures. If None, randomization is performed without loading additional data.
|
|
"""
|
|
|
|
def __init__(
|
|
self, scene_iter: Iterator[Scene], random_num: int = 1, strict_mode: bool = False, input_dir: str = None
|
|
):
|
|
super().__init__(scene_iter, random_num, strict_mode)
|
|
assert self.random_num > 0, "random_num must be greater than 0"
|
|
self.input_dir = input_dir
|
|
if self.input_dir is not None:
|
|
self.paths_names = os.listdir(self.input_dir)
|
|
self.random_num = len(self.paths_names)
|
|
|
|
def randomize_scene(self, scene) -> Scene:
|
|
if scene.plan_info is None:
|
|
path = None
|
|
if self.input_dir is not None:
|
|
path = os.path.join(self.input_dir, self.paths_names[self.cur_index])
|
|
if not scene.wf.randomization(path):
|
|
return None
|
|
else:
|
|
if not scene.wf.randomization_from_mem(scene.plan_info):
|
|
return None
|
|
return scene
|