Add: get_writer to LayerConfigManager

Author: Lars-Kraemer

src/vidata/config_manager.py

@@ -17,6 +17,14 @@
     SemanticSegmentationManager,
     TaskManager,
 )
+from vidata.writers import (
+    BaseWriter,
+    ImageStackWriter,
+    ImageWriter,
+    MultilabelStackedWriter,
+    MultilabelWriter,
+    SemSegWriter,
+)
 
 _VALID_SPLITS = ["train", "val", "test"]
 _IMAGE_LAYERS = {"image"}
@@ -30,6 +38,15 @@
     "image": ImageStackLoader,
     "multilabel": MultilabelStackedLoader,
 }
+_WRITER_MAPPING: dict[str, type[BaseWriter]] = {
+    "image": ImageWriter,
+    "semseg": SemSegWriter,
+    "multilabel": MultilabelWriter,
+}
+_STACKED_WRITER_MAPPING: dict[str, type[BaseWriter]] = {
+    "image": ImageStackWriter,
+    "multilabel": MultilabelStackedWriter,
+}
 
 
 class LayerConfigManager:
@@ -274,6 +291,25 @@ def data_loader(self) -> BaseLoader:
 
         return loader_cls(**args)
 
+    def data_writer(self) -> BaseWriter:
+        reg = _STACKED_WRITER_MAPPING if self.file_stack else _WRITER_MAPPING
+        try:
+            writer_cls = reg[self.type.lower()]
+        except KeyError as err:
+            raise ValueError(f"type {self.type} is not supported for layer {self.name}") from err
+
+        args = {
+            "ftype": self.file_type,
+            "backend": self.backend,
+        }
+
+        if self.type.lower() in _IMAGE_LAYERS:
+            args["channels"] = self.channels
+        elif self.type.lower() in _LABEL_LAYERS:
+            args["num_classes"] = self.classes
+
+        return writer_cls(**args)
+
     def task_manager(self) -> TaskManager:
         if self.type.lower() == "semseg":
             return SemanticSegmentationManager()

tests/config/test_config_manager.py

@@ -6,6 +6,7 @@
 from vidata.file_manager import FileManager
 from vidata.io import save_json
 from vidata.loaders import ImageLoader, MultilabelStackedLoader, SemSegLoader
+from vidata.writers import ImageWriter, MultilabelStackedWriter, SemSegWriter
 
 LEN_VAL = 4
 LEN_TRAIN = 6
@@ -141,13 +142,19 @@ def test_config_manager_splitfile(simple_config):
         assert layer.file_stack == l_conf.get("file_stack", False)
         assert layer.backend == l_conf["backend"]
 
-    for i, (ln, loader) in enumerate(
-        [("Images", ImageLoader), ("Labels", SemSegLoader), ("MLLabels", MultilabelStackedLoader)]
+    for i, (ln, loader, writer) in enumerate(
+        [
+            ("Images", ImageLoader, ImageWriter),
+            ("Labels", SemSegLoader, SemSegWriter),
+            ("MLLabels", MultilabelStackedLoader, MultilabelStackedWriter),
+        ]
     ):
         layer = cm.layer(ln)
 
         dl = layer.data_loader()
         assert isinstance(dl, loader)
+        df = layer.data_writer()
+        assert isinstance(df, writer)
 
         l_conf = layer.config()
         l_conf["pattern"] = simple_config["layers"][i]["pattern"]
@@ -188,13 +195,19 @@ def test_config_manager_splitfile_fold(simple_config):
         assert layer.file_stack == l_conf.get("file_stack", False)
         assert layer.backend == l_conf["backend"]
 
-    for i, (ln, loader) in enumerate(
-        [("Images", ImageLoader), ("Labels", SemSegLoader), ("MLLabels", MultilabelStackedLoader)]
+    for i, (ln, loader, writer) in enumerate(
+        [
+            ("Images", ImageLoader, ImageWriter),
+            ("Labels", SemSegLoader, SemSegWriter),
+            ("MLLabels", MultilabelStackedLoader, MultilabelStackedWriter),
+        ]
     ):
         layer = cm.layer(ln)
 
         dl = layer.data_loader()
         assert isinstance(dl, loader)
+        df = layer.data_writer()
+        assert isinstance(df, writer)
 
         l_conf = layer.config()
         l_conf["pattern"] = simple_config["layers"][i]["pattern"]

tox.ini

@@ -30,4 +30,4 @@ passenv =
     PYVISTA_OFF_SCREEN
 extras =
     testing
-commands = pytest -v --color=yes --cov=data_io --cov-report=xml
+commands = pytest -v --color=yes --cov=vidata --cov-report=xml