pipeline: Unifies pre- and postprocessing descr handling

Author: Tomaz-Vieira

bioimageio/core/proc_ops.py

@@ -660,72 +660,9 @@ def _apply(self, input: Tensor, stat: Stat) -> Tensor:
     ZeroMeanUnitVariance,
 ]
 
-
-def get_proc_class(proc_spec: ProcDescr):
-    if isinstance(proc_spec, (v0_4.BinarizeDescr, v0_5.BinarizeDescr)):
-        return Binarize
-    elif isinstance(proc_spec, (v0_4.ClipDescr, v0_5.ClipDescr)):
-        return Clip
-    elif isinstance(proc_spec, v0_5.EnsureDtypeDescr):
-        return EnsureDtype
-    elif isinstance(proc_spec, v0_5.FixedZeroMeanUnitVarianceDescr):
-        return FixedZeroMeanUnitVariance
-    elif isinstance(proc_spec, (v0_4.ScaleLinearDescr, v0_5.ScaleLinearDescr)):
-        return ScaleLinear
-    elif isinstance(
-        proc_spec, (v0_4.ScaleMeanVarianceDescr, v0_5.ScaleMeanVarianceDescr)
-    ):
-        return ScaleMeanVariance
-    elif isinstance(proc_spec, (v0_4.ScaleRangeDescr, v0_5.ScaleRangeDescr)):
-        return ScaleRange
-    elif isinstance(proc_spec, (v0_4.SigmoidDescr, v0_5.SigmoidDescr)):
-        return Sigmoid
-    elif (
-        isinstance(proc_spec, v0_4.ZeroMeanUnitVarianceDescr)
-        and proc_spec.kwargs.mode == "fixed"
-    ):
-        return FixedZeroMeanUnitVariance
-    elif isinstance(
-        proc_spec,
-        (v0_4.ZeroMeanUnitVarianceDescr, v0_5.ZeroMeanUnitVarianceDescr),
-    ):
-        return ZeroMeanUnitVariance
-    else:
-        assert_never(proc_spec)
-
-
-def preproc_v4_to_processing(
-    inp: v0_4.InputTensorDescr,
-    proc_spec: v0_4.PreprocessingDescr,
-) -> Processing:
-    member_id = TensorId(str(inp.name))
-    if isinstance(proc_spec, v0_4.BinarizeDescr):
-        return Binarize.from_proc_descr(proc_spec, member_id)
-    elif isinstance(proc_spec, v0_4.ClipDescr):
-        return Clip.from_proc_descr(proc_spec, member_id)
-    elif isinstance(proc_spec, v0_4.ScaleLinearDescr):
-        return ScaleLinear.from_proc_descr(proc_spec, member_id)
-    elif isinstance(proc_spec, v0_4.ScaleRangeDescr):
-        return ScaleRange.from_proc_descr(proc_spec, member_id)
-    elif isinstance(proc_spec, v0_4.SigmoidDescr):
-        return Sigmoid.from_proc_descr(proc_spec, member_id)
-    elif isinstance(proc_spec, v0_4.ZeroMeanUnitVarianceDescr):
-        if proc_spec.kwargs.mode == "fixed":
-            axes = inp.axes
-            v5_proc_spec = _convert_proc(proc_spec, axes)
-            assert isinstance(
-                v5_proc_spec, v0_5.FixedZeroMeanUnitVarianceDescr
-            )  # FIXME
-            return FixedZeroMeanUnitVariance.from_proc_descr(v5_proc_spec, member_id)
-        else:
-            return ZeroMeanUnitVariance.from_proc_descr(proc_spec, member_id)
-    else:
-        assert_never(proc_spec)
-
-
-def postproc_v4_to_processing(
-    inp: v0_4.OutputTensorDescr,
-    proc_spec: v0_4.PostprocessingDescr,
+def proc_descr_v4_to_op(
+    inp: Union[v0_4.InputTensorDescr, v0_4.OutputTensorDescr],
+    proc_spec: Union[v0_4.PreprocessingDescr, v0_4.PostprocessingDescr],
 ) -> Processing:
     member_id = TensorId(str(inp.name))
     if isinstance(proc_spec, v0_4.BinarizeDescr):
@@ -754,33 +691,9 @@ def postproc_v4_to_processing(
         assert_never(proc_spec)
 
 
-def preproc_v5_to_processing(
-    inp: v0_5.InputTensorDescr,
-    proc_spec: v0_5.PreprocessingDescr,
-) -> Processing:
-    if isinstance(proc_spec, v0_5.BinarizeDescr):
-        return Binarize.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.ClipDescr):
-        return Clip.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.ScaleLinearDescr):
-        return ScaleLinear.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.ScaleRangeDescr):
-        return ScaleRange.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.SigmoidDescr):
-        return Sigmoid.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.EnsureDtypeDescr):
-        return EnsureDtype.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.ZeroMeanUnitVarianceDescr):
-        return ZeroMeanUnitVariance.from_proc_descr(proc_spec, inp.id)
-    elif isinstance(proc_spec, v0_5.FixedZeroMeanUnitVarianceDescr):
-        return FixedZeroMeanUnitVariance.from_proc_descr(proc_spec, inp.id)
-    else:
-        assert_never(proc_spec)
-
-
-def postproc_v5_to_processing(
-    inp: v0_5.OutputTensorDescr,
-    proc_spec: v0_5.PostprocessingDescr,
+def proc_descr_v5_to_op(
+    inp: Union[v0_5.InputTensorDescr, v0_5.OutputTensorDescr],
+    proc_spec: Union[v0_5.PreprocessingDescr, v0_5.PostprocessingDescr],
 ) -> Processing:
     if isinstance(proc_spec, v0_5.BinarizeDescr):
         return Binarize.from_proc_descr(proc_spec, inp.id)

bioimageio/core/proc_setup.py

@@ -19,10 +19,8 @@
     EnsureDtype,
     Processing,
     UpdateStats,
-    postproc_v4_to_processing,
-    postproc_v5_to_processing,
-    preproc_v4_to_processing,
-    preproc_v5_to_processing,
+    proc_descr_v4_to_op,
+    proc_descr_v5_to_op,
 )
 from .sample import Sample
 from .stat_calculators import StatsCalculator
@@ -149,7 +147,7 @@ def _prepare_v4_preprocs(
             EnsureDtype(input=member_id, output=member_id, dtype=t_descr.data_type)
         )
         for proc_d in t_descr.preprocessing:
-            procs.append(preproc_v4_to_processing(t_descr, proc_d))
+            procs.append(proc_descr_v4_to_op(t_descr, proc_d))
     return procs
 
 
@@ -163,7 +161,7 @@ def _prepare_v4_postprocs(
             EnsureDtype(input=member_id, output=member_id, dtype=t_descr.data_type)
         )
         for proc_d in t_descr.postprocessing:
-            procs.append(postproc_v4_to_processing(t_descr, proc_d))
+            procs.append(proc_descr_v4_to_op(t_descr, proc_d))
     return procs
 
 
@@ -173,7 +171,7 @@ def _prepare_v5_preprocs(
     procs: List[Processing] = []
     for t_descr in tensor_descrs:
         for proc_d in t_descr.preprocessing:
-            procs.append(preproc_v5_to_processing(t_descr, proc_d))
+            procs.append(proc_descr_v5_to_op(t_descr, proc_d))
     return procs
 
 
@@ -183,7 +181,7 @@ def _prepare_v5_postprocs(
     procs: List[Processing] = []
     for t_descr in tensor_descrs:
         for proc_d in t_descr.postprocessing:
-            procs.append(postproc_v5_to_processing(t_descr, proc_d))
+            procs.append(proc_descr_v5_to_op(t_descr, proc_d))
     return procs