Task Manager Registry

Author: Lars-Kraemer

src/vidata/__init__.py

@@ -3,13 +3,22 @@
 # Registry imports
 from .config_manager import ConfigManager, LayerConfigManager
 from .file_manager import FileManager, FileManagerStacked
-from .registry import LOADER_REGISTRY, WRITER_REGISTRY, register_loader, register_writer
+from .registry import (
+    LOADER_REGISTRY,
+    TASK_REGISTRY,
+    WRITER_REGISTRY,
+    register_loader,
+    register_task,
+    register_writer,
+)
 
 __all__ = (
     "register_loader",
     "register_writer",
+    "register_task",
     "LOADER_REGISTRY",
     "WRITER_REGISTRY",
+    "TASK_REGISTRY",
     "FileManager",
     "FileManagerStacked",
     "ConfigManager",

src/vidata/config_manager.py

@@ -12,9 +12,8 @@
     MultilabelStackedLoader,
     SemSegLoader,
 )
+from vidata.registry import TASK_REGISTRY
 from vidata.task_manager import (
-    MultiLabelSegmentationManager,
-    SemanticSegmentationManager,
     TaskManager,
 )
 from vidata.writers import (
@@ -313,10 +312,8 @@ def data_writer(self) -> BaseWriter:
         return writer_cls(**args)
 
     def task_manager(self) -> TaskManager:
-        if self.type.lower() == "semseg":
-            return SemanticSegmentationManager()
-        elif self.type.lower() == "multilabel":
-            return MultiLabelSegmentationManager()
+        if self.type.lower() in TASK_REGISTRY:
+            return TASK_REGISTRY[self.type.lower()]
         else:
             raise ValueError(f"No Task manager defined for layer {self.name} and type {self.type}")
 

src/vidata/registry.py

@@ -13,6 +13,8 @@
     lambda: defaultdict(dict)
 )
 
+TASK_REGISTRY: dict[str, Any] = {}
+
 
 def register_loader(target: Target, *dtypes: str, backend: str = "default") -> Callable:
     """
@@ -52,6 +54,19 @@ def decorator(func: Callable) -> Callable:
     return decorator
 
 
+def register_task(name: str):
+    """Register a task class under a string identifier."""
+
+    def decorator(cls):
+        # if name in TASK_REGISTRY:
+        #     raise ValueError(f"Task '{name}' already registered.")
+        TASK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
 # --- Trigger backend imports ---
 # This must come LAST so the above decorators exist before data_io modules import them
 import vidata.io  # noqa
+import vidata.task_manager  # noqa

src/vidata/task_manager/image_manager.py

@@ -1,6 +1,9 @@
 import numpy as np
 
+from vidata.registry import register_task
 
+
+@register_task("image")
 class ImageManager:
     @staticmethod
     def random(size: tuple[int, ...], dtype="float") -> np.ndarray:

src/vidata/task_manager/multilabel_segmentation_manager.py

@@ -1,8 +1,10 @@
 import numpy as np
 
+from vidata.registry import register_task
 from vidata.task_manager.task_manager import TaskManager
 
 
+@register_task("multilabel")
 class MultiLabelSegmentationManager(TaskManager):
     @staticmethod
     def random(size: tuple[int, ...], num_classes: int) -> np.ndarray:

src/vidata/task_manager/semantic_segmentation_manager.py

@@ -1,8 +1,10 @@
 import numpy as np
 
+from vidata.registry import register_task
 from vidata.task_manager.task_manager import TaskManager
 
 
+@register_task("semseg")
 class SemanticSegmentationManager(TaskManager):
     @staticmethod
     def random(size: tuple[int, ...], num_classes: int) -> np.ndarray: