init commit

nimbus/scheduler/instructions.py

@@ -0,0 +1,115 @@
+from abc import abstractmethod
+
+from nimbus.components.dedump import dedumper_dict
+from nimbus.components.dump import dumper_dict
+from nimbus.components.load import layout_randomizer_dict, scene_loader_dict
+from nimbus.components.plan_with_render import plan_with_render_dict
+from nimbus.components.planner import seq_planner_dict
+from nimbus.components.render import renderer_dict
+from nimbus.components.store import writer_dict
+from nimbus.utils.types import ARGS, PLANNER, TYPE
+
+
+class Instruction:
+    def __init__(self, config):
+        self.config = config
+
+    @abstractmethod
+    def run(self, stage_input):
+        raise NotImplementedError()
+
+
+class LoadSceneInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.scene_iter = scene_loader_dict[self.config[TYPE]]
+
+    def run(self, stage_input):
+        pack_iter = pack_iter = stage_input.Args[0] if stage_input.Args is not None else None
+        return self.scene_iter(pack_iter=pack_iter, **self.config.get(ARGS, {}))
+
+
+class RandomizeLayoutInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.layout_randomlizer = layout_randomizer_dict[self.config[TYPE]]
+
+    def run(self, stage_input):
+        scene_iterator = stage_input.Args[0]
+        extend_scene_iterator = self.layout_randomlizer(scene_iterator, **self.config.get(ARGS, {}))
+        return extend_scene_iterator
+
+
+class PlanPathInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.seq_planner = seq_planner_dict[self.config[TYPE]]
+
+    def run(self, stage_input):
+        scene_iter = stage_input.Args[0]
+        planner_cfg = self.config[PLANNER] if PLANNER in self.config else None
+        return self.seq_planner(scene_iter, planner_cfg, **self.config.get(ARGS, {}))
+
+
+class RenderInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.renderer = renderer_dict[self.config[TYPE]]
+
+    def run(self, stage_input):
+        scene_seqs_iter = stage_input.Args[0]
+        obs_iter = self.renderer(scene_seqs_iter, **self.config.get(ARGS, {}))
+        return obs_iter
+
+
+class PlanWithRenderInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.plan_with_render = plan_with_render_dict[config[TYPE]]
+
+    def run(self, stage_input):
+        scene_iter = stage_input.Args[0]
+        plan_with_render_iter = self.plan_with_render(scene_iter, **self.config.get(ARGS, {}))
+        return plan_with_render_iter
+
+
+class StoreInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.writer = writer_dict[config[TYPE]]
+
+    def run(self, stage_input):
+        seqs_obs_iter = stage_input.Args[0]
+        store_iter = self.writer(seqs_obs_iter, **self.config.get(ARGS, {}))
+        return store_iter
+
+
+class DumpInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.dumper = dumper_dict[config[TYPE]]
+
+    def run(self, stage_input, output_queue=None):
+        seqs_obs_iter = stage_input.Args[0]
+        dump_iter = self.dumper(seqs_obs_iter, output_queue=output_queue, **self.config.get(ARGS, {}))
+        return dump_iter
+
+
+class DeDumpInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+        self.dedumper = dedumper_dict[config[TYPE]]
+
+    def run(self, stage_input, input_queue=None):
+        dump_iter = self.dedumper(input_queue=input_queue, **self.config.get(ARGS, {}))
+        return dump_iter
+
+
+class ComposeInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)
+
+
+class AnnotateDataInstruction(Instruction):
+    def __init__(self, config):
+        super().__init__(config)