small fixes and improve typing

bioimageio/core/_prediction_pipeline.py

@@ -46,8 +46,8 @@ def __init__(
         self._adapter: ModelAdapter = model
 
     def __call__(
-        self, *input_tensors: Tensor, **named_input_tensors: Tensor
-    ) -> List[Tensor]:
+        self, *input_tensors: Optional[Tensor], **named_input_tensors: Optional[Tensor]
+    ) -> List[Optional[Tensor]]:
         return self.forward(*input_tensors, **named_input_tensors)
 
     def __enter__(self):
@@ -60,7 +60,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore
 
     def predict(
         self, *input_tensors: Optional[Tensor], **named_input_tensors: Optional[Tensor]
-    ) -> List[Tensor]:
+    ) -> List[Optional[Tensor]]:
         """Predict input_tensor with the model without applying pre/postprocessing."""
         named_tensors = [
             named_input_tensors.get(str(k))
@@ -86,7 +86,12 @@ def forward_sample(self, input_sample: Sample) -> Sample:
             **{str(k): v for k, v in input_sample.data.items()}
         )
         prediction = Sample(
-            data=dict(zip(self.output_ids, prediction_tensors)), stat=input_sample.stat
+            data={
+                tid: t
+                for tid, t in zip(self.output_ids, prediction_tensors)
+                if t is not None
+            },
+            stat=input_sample.stat,
         )
         self.apply_postprocessing(prediction)
         return prediction

bioimageio/core/_resource_tests.py

@@ -127,17 +127,18 @@ def _test_model_inference(
             results = prediction_pipeline.forward(*inputs)
 
         if len(results) != len(expected):
-            error = (error or "") + (
-                f"Expected {len(expected)} outputs, but got {len(results)}"
-            )
+            error = f"Expected {len(expected)} outputs, but got {len(results)}"
+
         else:
             for res, exp in zip(results, expected):
+                if res is None:
+                    error = "Output tensors for test case may not be None"
+                    break
                 try:
                     np.testing.assert_array_almost_equal(res, exp, decimal=decimal)
                 except AssertionError as e:
-                    error = (
-                        error or ""
-                    ) + f"Output and expected output disagree:\n {e}"
+                    error = f"Output and expected output disagree:\n {e}"
+                    break
     except Exception as e:
         error = str(e)
         tb = traceback.format_tb(e.__traceback__)
@@ -238,11 +239,17 @@ def get_ns(n: int):
                 error: Optional[str] = None
                 results = prediction_pipeline.forward(*inputs)
                 if len(results) != len(exptected_output_shape):
-                    error = (error or "") + (
-                        f"Expected {len(exptected_output_shape)} outputs, but got {len(results)}"
+                    error = (
+                        f"Expected {len(exptected_output_shape)} outputs,"
+                        + f" but got {len(results)}"
                     )
+
                 else:
                     for res, exp in zip(results, exptected_output_shape):
+                        if res is None:
+                            error = "Output tensors may not be None for test case"
+                            break
+
                         diff: Dict[AxisId, int] = {}
                         for a, s in res.sizes.items():
                             if isinstance((e_aid := exp[AxisId(a)]), int):
@@ -254,10 +261,10 @@ def get_ns(n: int):
                                 diff[AxisId(a)] = s
                         if diff:
                             error = (
-                                (error or "")
-                                + f"(n={n}) Expected output shape {exp},"
+                                f"(n={n}) Expected output shape {exp},"
                                 + f" but got {res.sizes} ({diff})\n"
                             )
+                            break
 
                 model.validation_summary.add_detail(
                     ValidationDetail(

bioimageio/core/proc_ops.py

@@ -10,7 +10,6 @@
     Set,
     Tuple,
     Union,
-    cast,
 )
 
 import numpy as np
@@ -168,22 +167,29 @@ def __call__(self, sample: Sample) -> None:
 class Binarize(_SimpleOperator):
     """'output = tensor > threshold'."""
 
-    threshold: float
+    threshold: Union[float, Sequence[float]]
+    axis: Optional[AxisId] = None
 
     def _apply(self, input: Tensor, stat: Stat) -> xr.DataArray:
         return input > self.threshold
 
-    # @classmethod
-    # def from_descr(cls, descr: Union[v0_4.BinarizeDescr, v0_5.BinarizeDescr]):
-    #     return cls(threshold=descr.kwargs.threshold)
-
-    # def get_descr(self):
-    #     return v0_5.BinarizeDescr(kwargs=v0_5.BinarizeKwargs(threshold=self.threshold))
     @classmethod
     def from_proc_descr(
         cls, descr: Union[v0_4.BinarizeDescr, v0_5.BinarizeDescr], tensor_id: TensorId
     ) -> Self:
-        return cls(input=tensor_id, output=tensor_id, threshold=descr.kwargs.threshold)
+        if isinstance(descr.kwargs, (v0_4.BinarizeKwargs, v0_5.BinarizeKwargs)):
+            return cls(
+                input=tensor_id, output=tensor_id, threshold=descr.kwargs.threshold
+            )
+        elif isinstance(descr.kwargs, v0_5.BinarizeAlongAxisKwargs):
+            return cls(
+                input=tensor_id,
+                output=tensor_id,
+                threshold=descr.kwargs.threshold,
+                axis=descr.kwargs.axis,
+            )
+        else:
+            assert_never(descr.kwargs)
 
 
 @dataclass
@@ -224,7 +230,9 @@ def from_proc_descr(cls, descr: v0_5.EnsureDtypeDescr, tensor_id: TensorId):
 
     def get_descr(self):
         return v0_5.EnsureDtypeDescr(
-            kwargs=v0_5.EnsureDtypeKwargs(dtype=str(self.dtype))
+            kwargs=v0_5.EnsureDtypeKwargs(
+                dtype=str(self.dtype)  # pyright: ignore[reportArgumentType]
+            )
         )
 
     def _apply(self, input: Tensor, stat: Stat) -> Tensor:
@@ -242,25 +250,19 @@ class ScaleLinear(_SimpleOperator):
     def _apply(self, input: Tensor, stat: Stat) -> Tensor:
         return input * self.gain + self.offset
 
-    # @classmethod
-    # def from_descr(cls, descr: ScaleLinearDescr) -> Self:
-    #     ...
-
     @classmethod
     def from_proc_descr(
         cls,
         descr: Union[v0_4.ScaleLinearDescr, v0_5.ScaleLinearDescr],
         tensor_id: TensorId,
     ) -> Self:
         kwargs = descr.kwargs
-        if isinstance(kwargs, v0_5.ScaleLinearKwargs):
+        if isinstance(kwargs, v0_5.ScaleLinearAlongAxisKwargs):
             axis = kwargs.axis
-        elif kwargs.axes is not None:
-            raise NotImplementedError(
-                "ScaleLinear operator from v0_4.ScaleLinearDescr with axes"
-            )
-        else:
+        elif isinstance(kwargs, (v0_4.ScaleLinearKwargs, v0_5.ScaleLinearKwargs)):
             axis = None
+        else:
+            assert_never(kwargs)
 
         if axis:
             gain = xr.DataArray(np.atleast_1d(kwargs.gain), dims=axis)
@@ -535,29 +537,34 @@ def from_proc_descr(
         descr: v0_5.FixedZeroMeanUnitVarianceDescr,
         tensor_id: TensorId,
     ) -> Self:
+        if isinstance(descr.kwargs, v0_5.FixedZeroMeanUnitVarianceKwargs):
+            dims = None
+        elif isinstance(descr.kwargs, v0_5.FixedZeroMeanUnitVarianceAlongAxisKwargs):
+            dims = (descr.kwargs.axis,)
+        else:
+            assert_never(descr.kwargs)
+
         return cls(
             input=tensor_id,
             output=tensor_id,
-            mean=xr.DataArray(descr.kwargs.mean, dims=(descr.kwargs.axis,)),
-            std=xr.DataArray(descr.kwargs.std, dims=(descr.kwargs.axis,)),
+            mean=xr.DataArray(descr.kwargs.mean, dims=dims),
+            std=xr.DataArray(descr.kwargs.std, dims=dims),
         )
 
     def get_descr(self):
         if isinstance(self.mean, (int, float)):
             assert isinstance(self.std, (int, float))
-            axis = None
-            mean = self.mean
-            std = self.std
+            kwargs = v0_5.FixedZeroMeanUnitVarianceKwargs(mean=self.mean, std=self.std)
         else:
             assert isinstance(self.std, xr.DataArray)
             assert len(self.mean.dims) == 1
-            axis = AxisId(str(self.mean.dims[0]))
-            mean = tuple(self.mean)
-            std = tuple(self.std)
+            kwargs = v0_5.FixedZeroMeanUnitVarianceAlongAxisKwargs(
+                axis=AxisId(str(self.mean.dims[0])),
+                mean=list(self.mean),
+                std=list(self.std),
+            )
 
-        return v0_5.FixedZeroMeanUnitVarianceDescr(
-            kwargs=v0_5.FixedZeroMeanUnitVarianceKwargs(axis=axis, mean=mean, std=std)
-        )
+        return v0_5.FixedZeroMeanUnitVarianceDescr(kwargs=kwargs)
 
     def _apply(self, input: xr.DataArray, stat: Stat) -> xr.DataArray:
         return (input - self.mean) / (self.std + self.eps)

bioimageio/core/proc_setup.py

@@ -115,8 +115,8 @@ def prepare_procs(tensor_descrs: Sequence[TensorDescr]):
                     else t_descr.id
                 )
                 req = proc_class.from_proc_descr(
-                    proc_d, tensor_id
-                )  # pyright: ignore[reportArgumentType]
+                    proc_d, tensor_id  # pyright: ignore[reportArgumentType]
+                )
                 for m in req.required_measures:
                     if m.tensor_id in input_ids:
                         pre_measures.add(m)

bioimageio/core/weight_converter/keras/_tensorflow.py

@@ -1,3 +1,4 @@
+# type: ignore  # TODO: type
 import os
 import shutil
 from pathlib import Path
@@ -7,7 +8,7 @@
 try:
     import tensorflow.saved_model
 except Exception:
-    _tensorflow = None
+    tensorflow = None
 
 from bioimageio.spec._internal.io_utils import download
 from bioimageio.spec.model.v0_5 import ModelDescr
@@ -41,7 +42,9 @@ def _convert_tf1(
 ):
     try:
         # try to build the tf model with the keras import from tensorflow
-        from bioimageio.core.weight_converter.keras._tensorflow import keras  # type: ignore
+        from bioimageio.core.weight_converter.keras._tensorflow import (
+            keras,  # type: ignore
+        )
 
     except Exception:
         # if the above fails try to export with the standalone keras
@@ -50,20 +53,20 @@ def _convert_tf1(
     @no_type_check
     def build_tf_model():
         keras_model = keras.models.load_model(keras_weight_path)
-        assert _tensorflow is not None
-        builder = _tensorflow.saved_model.builder.SavedModelBuilder(output_path)
-        signature = _tensorflow.saved_model.signature_def_utils.predict_signature_def(
+        assert tensorflow is not None
+        builder = tensorflow.saved_model.builder.SavedModelBuilder(output_path)
+        signature = tensorflow.saved_model.signature_def_utils.predict_signature_def(
             inputs={input_name: keras_model.input},
             outputs={output_name: keras_model.output},
         )
 
         signature_def_map = {
-            _tensorflow.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
+            tensorflow.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
         }
 
         builder.add_meta_graph_and_variables(
             keras.backend.get_session(),
-            [_tensorflow.saved_model.tag_constants.SERVING],
+            [tensorflow.saved_model.tag_constants.SERVING],
             signature_def_map=signature_def_map,
         )
         builder.save()
@@ -107,8 +110,8 @@ def convert_weights_to_tensorflow_saved_model_bundle(
         model: The bioimageio model description
         output_path: where to save the tensorflow weights. This path must not exist yet.
     """
-    assert _tensorflow is not None
-    tf_major_ver = int(_tensorflow.__version__.split(".")[0])
+    assert tensorflow is not None
+    tf_major_ver = int(tensorflow.__version__.split(".")[0])
 
     if output_path.suffix == ".zip":
         output_path = output_path.with_suffix("")

bioimageio/core/weight_converter/torch/_onnx.py

@@ -1,3 +1,4 @@
+# type: ignore  # TODO: type
 import warnings
 from pathlib import Path
 from typing import Any, List, Sequence, cast

bioimageio/core/weight_converter/torch/_torchscript.py

@@ -1,3 +1,4 @@
+# type: ignore  # TODO: type
 from pathlib import Path
 from typing import List, Sequence, Union
 

bioimageio/core/weight_converter/torch/_utils.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 
 from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter
@@ -7,10 +9,14 @@
 
 # additional convenience for pytorch state dict, eventually we want this in python-bioimageio too
 # and for each weight format
-def load_torch_model(
-    node: "v0_4.PytorchStateDictWeightsDescr | v0_5.PytorchStateDictWeightsDescr",
+def load_torch_model(  # pyright: ignore[reportUnknownParameterType]
+    node: Union[v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr],
 ):
-    model = PytorchModelAdapter.get_network(node)
-    state = torch.load(download(node.source).path, map_location="cpu")
-    _ = model.load_state_dict(state)  # FIXME: check incompatible keys?
-    return model.eval()
+    model = (  # pyright: ignore[reportUnknownVariableType]
+        PytorchModelAdapter.get_network(node)
+    )
+    state = torch.load(  # pyright: ignore[reportUnknownVariableType]
+        download(node.source).path, map_location="cpu"
+    )
+    model.load_state_dict(state)  # FIXME: check incompatible keys?
+    return model.eval()  # pyright: ignore[reportUnknownVariableType]

tests/test_stat_measures.py

@@ -1,12 +1,12 @@
 from itertools import product
-from typing import Literal, Optional, Tuple
+from typing import Optional, Tuple
 
 import numpy as np
 import pytest
 import xarray as xr
 
 from bioimageio.core import stat_measures
-from bioimageio.core.common import AxisId, Sample, TensorId
+from bioimageio.core.common import AxisId, Sample, Tensor, TensorId
 from bioimageio.core.stat_calculators import (
     SamplePercentilesCalculator,
     get_measure_calculators,
@@ -29,7 +29,7 @@ def test_individual_normal_measure(
     measure = getattr(stat_measures, "Sample" + name.title())(
         axes=axes, tensor_id=data_id
     )
-    data = xr.DataArray(np.random.random((5, 6, 3)), dims=("x", "y", "c"))
+    data = Tensor(np.random.random((5, 6, 3)), dims=("x", "y", "c"))
 
     expected = getattr(data, name)(dim=axes)
     sample = Sample(data={data_id: data})
@@ -48,7 +48,7 @@ def test_individual_percentile_measure(axes: Optional[Tuple[AxisId, ...]]):
     calc = calcs[0]
     assert isinstance(calc, SamplePercentilesCalculator)
 
-    data = xr.DataArray(np.random.random((5, 6, 3)), dims=("x", "y", "c"))
+    data = Tensor(np.random.random((5, 6, 3)), dims=("x", "y", "c"))
     actual = calc.compute(Sample(data={tid: data}))
     for m in measures:
         expected = data.quantile(q=m.n / 100, dim=m.axes)