ExportImages step

src/htrflow_core/pipeline/steps.py

@@ -4,7 +4,7 @@
 from htrflow_core.dummies.dummy_models import simple_word_segmentation
 from htrflow_core.models.importer import all_models
 from htrflow_core.serialization import get_serializer
-from htrflow_core.utils.imgproc import binarize
+from htrflow_core.utils.imgproc import binarize, write
 from htrflow_core.volume.volume import Volume
 
 
@@ -32,7 +32,6 @@ def __str__(self):
 
 
 class Inference(PipelineStep):
-
     def __init__(self, model, generation_kwargs):
         self.model = model
         self.generation_kwargs = generation_kwargs
@@ -71,7 +70,6 @@ def run(self, volume):
 
 
 class WordSegmentation(PipelineStep):
-
     requires = [TextRecognition]
 
     def run(self, volume):
@@ -90,6 +88,18 @@ def run(self, volume):
         return volume
 
 
+class ExportImages(PipelineStep):
+    def __init__(self, dest):
+        self.dest = dest
+
+    def run(self, volume):
+        os.makedirs(self.dest, exist_ok=True)
+        for node in volume.traverse():
+            if node.image:
+                write(os.path.join(self.dest, id(node)), node.image)
+        return volume
+
+
 def auto_import(source) -> Volume:
     """Import volume from `source`
 
@@ -113,6 +123,7 @@ def auto_import(source) -> Volume:
 def all_subclasses(cls):
     return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in all_subclasses(c)])
 
+
 # Mapping class name -> class
 # Ex. {segmentation: `steps.Segmentation`}
 STEPS = {cls_.__name__.lower(): cls_ for cls_ in all_subclasses(PipelineStep)}