init commit

nimbus_extension/components/load/env_randomizer.py

@@ -0,0 +1,46 @@
+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