Redesign decoding module

Author: Donglai Wei

connectomics/decoding/__init__.py

@@ -1,22 +1,19 @@
-"""
-Decoding package for PyTorch Connectomics.
+"""Decoding package for PyTorch Connectomics."""
 
-This package provides post-processing functions to convert model predictions
-into final instance segmentation masks for various biological structures.
-
-Modules:
-    - segmentation: Mitochondria and organelle instance decoding
-    - synapse: Synaptic polarity instance decoding
-    - postprocess: General post-processing utilities
-    - utils: Shared utility functions
-    - auto_tuning: Hyperparameter optimization for post-processing
-
-Import patterns:
-    from connectomics.decoding import decode_binary_watershed, decode_binary_contour_watershed
-    from connectomics.decoding import polarity2instance
-    from connectomics.decoding import stitch_3d, watershed_split
-    from connectomics.decoding import optimize_threshold, SkeletonMetrics
-"""
+from .base import DecodeStep
+from .pipeline import (
+    apply_decode_mode,
+    apply_decode_pipeline,
+    normalize_decode_modes,
+    resolve_decode_modes_from_cfg,
+)
+from .registry import (
+    DecoderRegistry,
+    get_decoder,
+    list_decoders,
+    register_builtin_decoders,
+    register_decoder,
+)
 
 from .auto_tuning import (
     SkeletonMetrics,
@@ -53,7 +50,19 @@
     remove_small_instances,
 )
 
+register_builtin_decoders()
+
 __all__ = [
+    # Registry / pipeline
+    "DecodeStep",
+    "DecoderRegistry",
+    "register_decoder",
+    "get_decoder",
+    "list_decoders",
+    "normalize_decode_modes",
+    "apply_decode_pipeline",
+    "resolve_decode_modes_from_cfg",
+    "apply_decode_mode",
     # Segmentation decoding
     "decode_instance_binary_contour_distance",
     "decode_affinity_cc",

connectomics/decoding/base.py

@@ -0,0 +1,23 @@
+"""Core decoding type definitions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Protocol
+
+import numpy as np
+
+
+class DecodeFunction(Protocol):
+    """Callable decoder signature used by the registry."""
+
+    def __call__(self, predictions: np.ndarray, **kwargs: Any) -> np.ndarray: ...
+
+
+@dataclass
+class DecodeStep:
+    """Single step in a decoding pipeline."""
+
+    name: str
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+

connectomics/decoding/pipeline.py

@@ -0,0 +1,136 @@
+"""Decoding pipeline helpers."""
+
+from __future__ import annotations
+
+from typing import Any, Iterable, List, Sequence, Tuple
+
+import numpy as np
+
+from .base import DecodeStep
+from .registry import DecoderRegistry, DEFAULT_DECODER_REGISTRY, register_builtin_decoders
+
+
+def _coerce_kwargs(kwargs: Any) -> dict:
+    if kwargs is None:
+        return {}
+    if hasattr(kwargs, "items"):
+        return dict(kwargs)
+    raise TypeError(
+        f"Decode kwargs must be a mapping, got {type(kwargs).__name__}"
+    )
+
+
+def normalize_decode_modes(decode_modes: Iterable[Any]) -> List[DecodeStep]:
+    """Normalize decode configuration entries into DecodeStep objects."""
+    steps: List[DecodeStep] = []
+    for mode in decode_modes:
+        if isinstance(mode, DecodeStep):
+            steps.append(DecodeStep(name=mode.name, kwargs=_coerce_kwargs(mode.kwargs)))
+            continue
+
+        if hasattr(mode, "name"):
+            name = mode.name
+            kwargs = _coerce_kwargs(getattr(mode, "kwargs", {}))
+            steps.append(DecodeStep(name=name, kwargs=kwargs))
+            continue
+
+        if isinstance(mode, dict):
+            name = mode.get("name")
+            kwargs = _coerce_kwargs(mode.get("kwargs", {}))
+            steps.append(DecodeStep(name=name, kwargs=kwargs))
+            continue
+
+        raise TypeError(f"Unsupported decode mode type: {type(mode).__name__}")
+
+    for step in steps:
+        if not step.name:
+            raise ValueError("Decode step is missing required field 'name'.")
+
+    return steps
+
+
+def _prepare_batched_input(data: np.ndarray) -> Tuple[np.ndarray, int]:
+    arr = np.asarray(data)
+    if arr.ndim == 5:
+        return arr, arr.shape[0]
+    if arr.ndim == 4:
+        return arr[np.newaxis, ...], 1
+    if arr.ndim == 3:
+        return arr[np.newaxis, np.newaxis, ...], 1
+    if arr.ndim == 2:
+        return arr[np.newaxis, np.newaxis, np.newaxis, ...], 1
+    raise ValueError(
+        f"Expected input with 2-5 dimensions, got shape {arr.shape}."
+    )
+
+
+def apply_decode_pipeline(
+    data: np.ndarray,
+    decode_modes: Sequence[Any] | None,
+    registry: DecoderRegistry | None = None,
+) -> np.ndarray:
+    """Apply configured decode steps to prediction data."""
+    if not decode_modes:
+        return data
+
+    register_builtin_decoders()
+    registry = registry or DEFAULT_DECODER_REGISTRY
+    steps = normalize_decode_modes(decode_modes)
+    batched, batch_size = _prepare_batched_input(data)
+
+    results: List[np.ndarray] = []
+    for batch_idx in range(batch_size):
+        sample = batched[batch_idx]
+        for step in steps:
+            try:
+                decoder = registry.get(step.name)
+            except KeyError as exc:
+                available = ", ".join(registry.available())
+                raise ValueError(
+                    f"Unknown decode function '{step.name}'. "
+                    f"Available functions: [{available}]."
+                ) from exc
+
+            try:
+                sample = decoder(sample, **step.kwargs)
+            except Exception as exc:
+                raise RuntimeError(
+                    f"Error applying decode function '{step.name}': {exc}"
+                ) from exc
+
+        results.append(sample)
+
+    if len(results) == 1:
+        return results[0]
+    return np.stack(results, axis=0)
+
+
+def resolve_decode_modes_from_cfg(cfg: Any) -> Sequence[Any] | None:
+    """Resolve decode mode list from config.
+
+    Priority:
+    1. ``cfg.test.decoding``
+    2. ``cfg.inference.decoding``
+    """
+    if hasattr(cfg, "test") and cfg.test and hasattr(cfg.test, "decoding"):
+        return cfg.test.decoding
+    if hasattr(cfg, "inference") and hasattr(cfg.inference, "decoding"):
+        return cfg.inference.decoding
+    return None
+
+
+def apply_decode_mode(cfg: Any, data: np.ndarray, *, verbose: bool = True) -> np.ndarray:
+    """Apply decode pipeline resolved from ``test.decoding`` or ``inference.decoding``."""
+    decode_modes = resolve_decode_modes_from_cfg(cfg)
+    if not decode_modes:
+        if verbose:
+            print("  No decoding configuration found (test.decoding or inference.decoding)")
+        return data
+
+    if verbose:
+        source = "test.decoding"
+        if not (hasattr(cfg, "test") and cfg.test and hasattr(cfg.test, "decoding")):
+            source = "inference.decoding"
+        print(f"  Using {source}: {decode_modes}")
+
+    return apply_decode_pipeline(data, decode_modes)

connectomics/decoding/registry.py

@@ -0,0 +1,80 @@
+"""Decoder registry for configurable decode pipelines."""
+
+from __future__ import annotations
+
+from typing import Dict, List
+
+from .base import DecodeFunction
+
+
+class DecoderRegistry:
+    """Name -> decoder function registry."""
+
+    def __init__(self) -> None:
+        self._decoders: Dict[str, DecodeFunction] = {}
+
+    def register(self, name: str, fn: DecodeFunction, *, overwrite: bool = False) -> None:
+        """Register a decoder function."""
+        if not name or not isinstance(name, str):
+            raise ValueError("Decoder name must be a non-empty string.")
+        if not callable(fn):
+            raise TypeError(f"Decoder '{name}' must be callable.")
+
+        existing = self._decoders.get(name)
+        if existing is not None and not overwrite and existing is not fn:
+            raise ValueError(
+                f"Decoder '{name}' already registered. Use overwrite=True to replace it."
+            )
+        self._decoders[name] = fn
+
+    def get(self, name: str) -> DecodeFunction:
+        """Get decoder function by name."""
+        try:
+            return self._decoders[name]
+        except KeyError as exc:
+            available = ", ".join(sorted(self._decoders))
+            raise KeyError(
+                f"Unknown decoder '{name}'. Available decoders: [{available}]"
+            ) from exc
+
+    def available(self) -> List[str]:
+        """Return registered decoder names."""
+        return sorted(self._decoders.keys())
+
+
+DEFAULT_DECODER_REGISTRY = DecoderRegistry()
+
+
+def register_decoder(name: str, fn: DecodeFunction, *, overwrite: bool = False) -> None:
+    """Register decoder in the default registry."""
+    DEFAULT_DECODER_REGISTRY.register(name, fn, overwrite=overwrite)
+
+
+def get_decoder(name: str) -> DecodeFunction:
+    """Get decoder from the default registry."""
+    return DEFAULT_DECODER_REGISTRY.get(name)
+
+
+def list_decoders() -> List[str]:
+    """List names in the default registry."""
+    return DEFAULT_DECODER_REGISTRY.available()
+
+
+def register_builtin_decoders() -> None:
+    """Populate registry with built-in decoders."""
+    from .segmentation import (
+        decode_affinity_cc,
+        decode_distance_watershed,
+        decode_instance_binary_contour_distance,
+    )
+    from .synapse import polarity2instance
+
+    register_decoder(
+        "decode_instance_binary_contour_distance",
+        decode_instance_binary_contour_distance,
+        overwrite=True,
+    )
+    register_decoder("decode_affinity_cc", decode_affinity_cc, overwrite=True)
+    register_decoder("decode_distance_watershed", decode_distance_watershed, overwrite=True)
+    register_decoder("polarity2instance", polarity2instance, overwrite=True)
+

connectomics/inference/__init__.py

@@ -1,21 +1,15 @@
 """Inference utilities package."""
 
 from .manager import InferenceManager
-from .io import (
-    apply_save_prediction_transform,
-    apply_postprocessing,
-    apply_decode_mode,
-    resolve_output_filenames,
-    write_outputs,
-)
+from .output import resolve_output_filenames, write_outputs
+from .postprocessing import apply_save_prediction_transform, apply_postprocessing
 from .sliding import build_sliding_inferer, resolve_inferer_roi_size, resolve_inferer_overlap
 from .tta import TTAPredictor
 
 __all__ = [
     "InferenceManager",
     "apply_save_prediction_transform",
     "apply_postprocessing",
-    "apply_decode_mode",
     "resolve_output_filenames",
     "write_outputs",
     "build_sliding_inferer",