setup

battle_logger/pipeline.py

@@ -1,4 +1,5 @@
-from typing import List
+import os
+from typing import List, Optional
 
 from algo_ops.ops.op import Op
 from algo_ops.pipeline.pipeline import Pipeline
@@ -31,6 +32,7 @@ def get_arch() -> List[Op]:
                 output_type=OutputType.TEXTBOX,
                 text_pipeline=None,
                 autosave_output_img_path=None,
+                store_intermediate_images=False,
             ),
             BattleLoggerOp(),
         ]
@@ -53,5 +55,25 @@ def set_output_paths(
         ffmpeg_op.image_out_path = image_out_path
         ocr_pipeline_op.ocr_op.autosave_output_img_path = autosave_output_img_path
 
+    def save_output(self, out_path: str = ".", basename: Optional[str] = None) -> None:
+        """
+        Saves pipeline Op outputs to file.
+
+        param out_path: Path to where output should go
+        param basename: Basename of output file
+        """
+        os.makedirs(out_path, exist_ok=True)
+        for i, op_name in enumerate(self.ops.keys()):
+            op = self.ops[op_name]
+            assert isinstance(op, Op)
+            op_pipeline_name = self._pipeline_op_name(op=op)
+            if i == 0:
+                op.save_input(out_path=out_path, basename=op_pipeline_name)
+            # This is hack to get around dropping of input/output images in OCRPipeline
+            try:
+                op.save_output(out_path=out_path, basename=op_pipeline_name)
+            except ValueError:
+                pass
+
     def __init__(self):
         super().__init__(ops=self.get_arch())

setup.py

@@ -7,7 +7,7 @@
 
 setuptools.setup(
     name="pokemon-go-video-analysis",
-    version="0.0.0.0.5.2",
+    version="0.0.0.0.5.3",
     author="Prateek Tandon",
     author_email="[email protected]",
     description="Pokemon Go Video Analysis",