add AssertDtype Processing

Author: FynnBe

bioimageio/core/prediction_pipeline/_combined_processing.py

@@ -1,8 +1,8 @@
 import dataclasses
-from typing import Dict, Iterable, List, NamedTuple, Optional, Sequence, TypedDict, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 from bioimageio.core.resource_io import nodes
-from ._processing import EnsureDtype, KNOWN_PROCESSING, Processing, TensorName
+from ._processing import AssertDtype, EnsureDtype, KNOWN_PROCESSING, Processing, TensorName
 from ._utils import ComputedMeasures, PER_DATASET, PER_SAMPLE, RequiredMeasures, Sample
 
 try:
@@ -11,11 +11,19 @@
     from typing_extensions import Literal  # type: ignore
 
 
+@dataclasses.dataclass
+class Processing:
+    name: str
+    kwargs: Dict[str, Any]
+
+
 @dataclasses.dataclass
 class TensorProcessingInfo:
-    processing_steps: Union[List[nodes.Preprocessing], List[nodes.Postprocessing]]
-    data_type_before: Optional[str] = None
-    data_type_after: Optional[str] = None
+    processing_steps: List[Processing]
+    assert_dtype_before: Optional[Union[str, Sequence[str]]] = None  # throw AssertionError if data type doesn't match
+    ensure_dtype_before: Optional[str] = None  # cast data type if needed
+    assert_dtype_after: Optional[Union[str, Sequence[str]]] = None  # throw AssertionError if data type doesn't match
+    ensure_dtype_after: Optional[str] = None  # throw AssertionError if data type doesn't match
 
 
 class CombinedProcessing:
@@ -26,16 +34,22 @@ def __init__(self, combine_tensors: Dict[TensorName, TensorProcessingInfo]):
 
         # ensure all tensors have correct data type before any processing
         for tensor_name, info in combine_tensors.items():
-            if info.data_type_before is not None:
-                self._procs.append(EnsureDtype(tensor_name=tensor_name, dtype=info.data_type_before))
+            if info.assert_dtype_before is not None:
+                self._procs.append(AssertDtype(tensor_name=tensor_name, dtype=info.assert_dtype_before))
+
+            if info.ensure_dtype_before is not None:
+                self._procs.append(EnsureDtype(tensor_name=tensor_name, dtype=info.ensure_dtype_before))
 
         for tensor_name, info in combine_tensors.items():
             for step in info.processing_steps:
                 self._procs.append(known[step.name](tensor_name=tensor_name, **step.kwargs))
 
+            if info.assert_dtype_after is not None:
+                self._procs.append(AssertDtype(tensor_name=tensor_name, dtype=info.assert_dtype_after))
+
             # ensure tensor has correct data type right after its processing
-            if info.data_type_after is not None:
-                self._procs.append(EnsureDtype(tensor_name=tensor_name, dtype=info.data_type_after))
+            if info.ensure_dtype_after is not None:
+                self._procs.append(EnsureDtype(tensor_name=tensor_name, dtype=info.ensure_dtype_after))
 
         self.required_measures: RequiredMeasures = self._collect_required_measures(self._procs)
         self.tensor_names = list(combine_tensors)
@@ -50,10 +64,15 @@ def from_tensor_specs(cls, tensor_specs: List[Union[nodes.InputTensor, nodes.Out
             # todo: cast dtype for inputs before preprocessing? or check dtype?
             assert ts.name not in combine_tensors
             if isinstance(ts, nodes.InputTensor):
-                # todo: move preprocessing ensure_dtype here as data_type_after
-                combine_tensors[ts.name] = TensorProcessingInfo(ts.preprocessing)
+                # todo: assert nodes.InputTensor.dtype with assert_dtype_before?
+                combine_tensors[ts.name] = TensorProcessingInfo(
+                    [Processing(p.name, kwargs=p.kwargs) for p in ts.preprocessing]
+                )
             elif isinstance(ts, nodes.OutputTensor):
-                combine_tensors[ts.name] = TensorProcessingInfo(ts.postprocessing, None, ts.data_type)
+                combine_tensors[ts.name] = TensorProcessingInfo(
+                    [Processing(p.name, kwargs=p.kwargs) for p in ts.postprocessing],
+                    ensure_dtype_after=ts.data_type,
+                )
             else:
                 raise NotImplementedError(type(ts))
 

bioimageio/core/prediction_pipeline/_processing.py

@@ -2,9 +2,11 @@
 see https://github.com/bioimage-io/spec-bioimage-io/blob/gh-pages/preprocessing_spec_latest.md
 and https://github.com/bioimage-io/spec-bioimage-io/blob/gh-pages/postprocessing_spec_latest.md
 """
-from dataclasses import dataclass, field, fields
-from typing import Mapping, Optional, Sequence, Type, Union
+import numbers
+from dataclasses import InitVar, dataclass, field, fields
+from typing import List, Mapping, Optional, Sequence, Tuple, Type, Union
 
+import numpy
 import numpy as np
 import xarray as xr
 
@@ -105,6 +107,26 @@ def ensure_dtype(tensor: xr.DataArray, *, dtype) -> xr.DataArray:
 #
 
 
+@dataclass
+class AssertDtype(Processing):
+    """Helper Processing to assert dtype."""
+
+    dtype: Union[str, Sequence[str]] = MISSING
+    assert_with: Tuple[Type[numpy.dtype], ...] = field(init=False)
+
+    def __post_init__(self):
+        if isinstance(self.dtype, str):
+            dtype = [self.dtype]
+        else:
+            dtype = self.dtype
+
+        self.assert_with = tuple(type(numpy.dtype(dt)) for dt in dtype)
+
+    def apply(self, tensor: xr.DataArray) -> xr.DataArray:
+        assert isinstance(tensor.dtype, self.assert_with)
+        return tensor
+
+
 @dataclass
 class Binarize(Processing):
     """'output = tensor > threshold' (note: returns float array)."""

tests/prediction_pipeline/test_processing.py

@@ -1,6 +1,8 @@
 import dataclasses
 
+import numpy as np
 import pytest
+import xarray as xr
 
 from bioimageio.core.prediction_pipeline._processing import KNOWN_PROCESSING
 from bioimageio.core.prediction_pipeline._utils import FIXED
@@ -11,6 +13,20 @@
     from typing_extensions import get_args  # type: ignore
 
 
+def test_assert_dtype():
+    from bioimageio.core.prediction_pipeline._processing import AssertDtype
+
+    proc = AssertDtype("test_tensor", dtype="uint8")
+    tensor = xr.DataArray(np.zeros((1,), dtype="uint8"), dims=tuple("c"))
+    out = proc(tensor)
+    assert out is tensor
+
+    tensor = tensor.astype("uint16")
+    with pytest.raises(AssertionError):
+        out = proc(tensor)
+        assert out is tensor
+
+
 @pytest.mark.parametrize(
     "proc",
     list(KNOWN_PROCESSING["pre"].values()) + list(KNOWN_PROCESSING["post"].values()),