Change: Group commands

Pipfile

@@ -20,8 +20,8 @@ fire = "*"
 [scripts]
 fix = "autopep8 -ivr ."
 lint = "flake8 ."
-train = "python -m trainer train"
-preview = "python -m trainer preview"
+train = "python -m trainer trainer train"
+preview = "python -m trainer trainer preview"
 
 [requires]
 python_version = "3.7"

trainer/__main__.py

@@ -1,6 +1,6 @@
 import fire
 
-from trainer import cli
+from trainer.cli import Trainer, Utility
 
 if __name__ == '__main__':
-    fire.Fire(cli)
+    fire.Fire({'trainer': Trainer, 'util': Utility})

trainer/cli.py

@@ -2,32 +2,53 @@
 
 import trainer
 
-def train(motion, output, robot, timestep=0.0165/4, frame_skip=4, chunk_length=3, num_iteration=1000, num_chunk=50, weight_factor=0.01):
-    m = flom.load(motion)
-    trained = trainer.train(m, robot, timestep, frame_skip, chunk_length, num_iteration, num_chunk, weight_factor)
-    trained.dump(output)
+import dataclasses
+
+
+@dataclasses.dataclass
+class Trainer:
+    motion: dataclasses.InitVar[str]
+    robot: str
+    timestep: float = 0.0165/4
+    frame_skip: int = 4
+
+    input_motion: flom.Motion = dataclasses.field(init=False)
+
+    def __post_init__(self, motion):
+        self.input_motion = flom.load(motion)
+
+    def train(self, output, chunk_length=3, num_iteration=1000, num_chunk=50, weight_factor=0.01):
+        trained = trainer.train(self.input_motion, self.robot, self.timestep,
+                                self.frame_skip, chunk_length, num_iteration, num_chunk, weight_factor)
+        trained.dump(output)
+
+    def preview(self):
+        trainer.preview(self.input_motion, self.robot, self.timestep, self.frame_skip)
+
 
 # TODO: Move these utilities to the new package
+@dataclasses.dataclass
+class Utility:
+    motion: dataclasses.InitVar[str]
+    output: str
 
-def preview(motion, robot, timestep=0.0165/4, frame_skip=4):
-    m = flom.load(motion)
-    trainer.preview(m, robot, timestep, frame_skip)
+    input_motion: flom.Motion = dataclasses.field(init=False)
 
-def add_noise(motion, output, random=0.1):
-    m = flom.load(motion)
-    trainer.utils.add_noise(m, random)
-    m.dump(output)
+    def __post_init__(self, motion):
+        self.input_motion = flom.load(motion)
 
-def plot(motion, output, fps=0.01, loop=1):
-    from trainer import plot
-    import matplotlib.pyplot as plt
+    def add_noise(self, random=0.1):
+        trainer.utils.add_noise(self.input_motion, random)
+        self.input_motion.dump(self.output)
 
-    m = flom.load(motion)
+    def plot(self, fps=0.01, loop=1):
+        from trainer import plot
+        import matplotlib.pyplot as plt
 
-    fig = plt.figure()
-    ax = fig.add_subplot(1, 1, 1)
+        fig = plt.figure()
+        ax = fig.add_subplot(1, 1, 1)
 
-    plot.plot_frames(m, ax, loop, fps)
+        plot.plot_frames(self.input_motion, ax, loop, fps)
 
-    ax.legend()
-    plt.savefig(output)
+        ax.legend()
+        plt.savefig(self.output)