116 lines
3.7 KiB
Python
116 lines
3.7 KiB
Python
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)
|