bioimageio.core.proc_setup

  1from typing import (
  2    Callable,
  3    Iterable,
  4    List,
  5    Mapping,
  6    NamedTuple,
  7    Optional,
  8    Sequence,
  9    Set,
 10    Union,
 11)
 12
 13from typing_extensions import assert_never
 14
 15from bioimageio.core.digest_spec import get_member_id
 16from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5
 17
 18from .proc_ops import (
 19    AddKnownDatasetStats,
 20    EnsureDtype,
 21    Processing,
 22    UpdateStats,
 23    get_proc,
 24)
 25from .sample import Sample
 26from .stat_calculators import StatsCalculator
 27from .stat_measures import (
 28    DatasetMeasure,
 29    DatasetMeasureBase,
 30    Measure,
 31    MeasureValue,
 32    SampleMeasure,
 33    SampleMeasureBase,
 34)
 35
 36TensorDescr = Union[
 37    v0_4.InputTensorDescr,
 38    v0_4.OutputTensorDescr,
 39    v0_5.InputTensorDescr,
 40    v0_5.OutputTensorDescr,
 41]
 42
 43
 44class PreAndPostprocessing(NamedTuple):
 45    pre: List[Processing]
 46    post: List[Processing]
 47
 48
 49class _ProcessingCallables(NamedTuple):
 50    pre: Callable[[Sample], None]
 51    post: Callable[[Sample], None]
 52
 53
 54class _SetupProcessing(NamedTuple):
 55    pre: List[Processing]
 56    post: List[Processing]
 57    pre_measures: Set[Measure]
 58    post_measures: Set[Measure]
 59
 60
 61class _ApplyProcs:
 62    def __init__(self, procs: Sequence[Processing]):
 63        super().__init__()
 64        self._procs = procs
 65
 66    def __call__(self, sample: Sample) -> None:
 67        for op in self._procs:
 68            op(sample)
 69
 70
 71def get_pre_and_postprocessing(
 72    model: AnyModelDescr,
 73    *,
 74    dataset_for_initial_statistics: Iterable[Sample],
 75    keep_updating_initial_dataset_stats: bool = False,
 76    fixed_dataset_stats: Optional[Mapping[DatasetMeasure, MeasureValue]] = None,
 77) -> _ProcessingCallables:
 78    """Creates callables to apply pre- and postprocessing in-place to a sample"""
 79
 80    setup = setup_pre_and_postprocessing(
 81        model=model,
 82        dataset_for_initial_statistics=dataset_for_initial_statistics,
 83        keep_updating_initial_dataset_stats=keep_updating_initial_dataset_stats,
 84        fixed_dataset_stats=fixed_dataset_stats,
 85    )
 86    return _ProcessingCallables(_ApplyProcs(setup.pre), _ApplyProcs(setup.post))
 87
 88
 89def setup_pre_and_postprocessing(
 90    model: AnyModelDescr,
 91    dataset_for_initial_statistics: Iterable[Sample],
 92    keep_updating_initial_dataset_stats: bool = False,
 93    fixed_dataset_stats: Optional[Mapping[DatasetMeasure, MeasureValue]] = None,
 94) -> PreAndPostprocessing:
 95    """
 96    Get pre- and postprocessing operators for a `model` description.
 97    Used in `bioimageio.core.create_prediction_pipeline"""
 98    prep, post, prep_meas, post_meas = _prepare_setup_pre_and_postprocessing(model)
 99
100    missing_dataset_stats = {
101        m
102        for m in prep_meas | post_meas
103        if fixed_dataset_stats is None or m not in fixed_dataset_stats
104    }
105    if missing_dataset_stats:
106        initial_stats_calc = StatsCalculator(missing_dataset_stats)
107        for sample in dataset_for_initial_statistics:
108            initial_stats_calc.update(sample)
109
110        initial_stats = initial_stats_calc.finalize()
111    else:
112        initial_stats = {}
113
114    prep.insert(
115        0,
116        UpdateStats(
117            StatsCalculator(prep_meas, initial_stats),
118            keep_updating_initial_dataset_stats=keep_updating_initial_dataset_stats,
119        ),
120    )
121    if post_meas:
122        post.insert(
123            0,
124            UpdateStats(
125                StatsCalculator(post_meas, initial_stats),
126                keep_updating_initial_dataset_stats=keep_updating_initial_dataset_stats,
127            ),
128        )
129
130    if fixed_dataset_stats:
131        prep.insert(0, AddKnownDatasetStats(fixed_dataset_stats))
132        post.insert(0, AddKnownDatasetStats(fixed_dataset_stats))
133
134    return PreAndPostprocessing(prep, post)
135
136
137class RequiredMeasures(NamedTuple):
138    pre: Set[Measure]
139    post: Set[Measure]
140
141
142class RequiredDatasetMeasures(NamedTuple):
143    pre: Set[DatasetMeasure]
144    post: Set[DatasetMeasure]
145
146
147class RequiredSampleMeasures(NamedTuple):
148    pre: Set[SampleMeasure]
149    post: Set[SampleMeasure]
150
151
152def get_requried_measures(model: AnyModelDescr) -> RequiredMeasures:
153    s = _prepare_setup_pre_and_postprocessing(model)
154    return RequiredMeasures(s.pre_measures, s.post_measures)
155
156
157def get_required_dataset_measures(model: AnyModelDescr) -> RequiredDatasetMeasures:
158    s = _prepare_setup_pre_and_postprocessing(model)
159    return RequiredDatasetMeasures(
160        {m for m in s.pre_measures if isinstance(m, DatasetMeasureBase)},
161        {m for m in s.post_measures if isinstance(m, DatasetMeasureBase)},
162    )
163
164
165def get_requried_sample_measures(model: AnyModelDescr) -> RequiredSampleMeasures:
166    s = _prepare_setup_pre_and_postprocessing(model)
167    return RequiredSampleMeasures(
168        {m for m in s.pre_measures if isinstance(m, SampleMeasureBase)},
169        {m for m in s.post_measures if isinstance(m, SampleMeasureBase)},
170    )
171
172
173def _prepare_procs(
174    tensor_descrs: Union[
175        Sequence[v0_4.InputTensorDescr],
176        Sequence[v0_5.InputTensorDescr],
177        Sequence[v0_4.OutputTensorDescr],
178        Sequence[v0_5.OutputTensorDescr],
179    ],
180) -> List[Processing]:
181    procs: List[Processing] = []
182    for t_descr in tensor_descrs:
183        if isinstance(t_descr, (v0_4.InputTensorDescr, v0_4.OutputTensorDescr)):
184            member_id = get_member_id(t_descr)
185            procs.append(
186                EnsureDtype(input=member_id, output=member_id, dtype=t_descr.data_type)
187            )
188
189        if isinstance(t_descr, (v0_4.InputTensorDescr, v0_5.InputTensorDescr)):
190            for proc_d in t_descr.preprocessing:
191                procs.append(get_proc(proc_d, t_descr))
192        elif isinstance(t_descr, (v0_4.OutputTensorDescr, v0_5.OutputTensorDescr)):
193            for proc_d in t_descr.postprocessing:
194                procs.append(get_proc(proc_d, t_descr))
195        else:
196            assert_never(t_descr)
197
198        if isinstance(
199            t_descr,
200            (v0_4.InputTensorDescr, (v0_4.InputTensorDescr, v0_4.OutputTensorDescr)),
201        ):
202            if len(procs) == 1:
203                # remove initial ensure_dtype if there are no other proccessing steps
204                assert isinstance(procs[0], EnsureDtype)
205                procs = []
206
207            # ensure 0.4 models get float32 input
208            # which has been the implicit assumption for 0.4
209            member_id = get_member_id(t_descr)
210            procs.append(
211                EnsureDtype(input=member_id, output=member_id, dtype="float32")
212            )
213
214    return procs
215
216
217def _prepare_setup_pre_and_postprocessing(model: AnyModelDescr) -> _SetupProcessing:
218    if isinstance(model, v0_4.ModelDescr):
219        pre = _prepare_procs(model.inputs)
220        post = _prepare_procs(model.outputs)
221    elif isinstance(model, v0_5.ModelDescr):
222        pre = _prepare_procs(model.inputs)
223        post = _prepare_procs(model.outputs)
224    else:
225        assert_never(model)
226
227    return _SetupProcessing(
228        pre=pre,
229        post=post,
230        pre_measures={m for proc in pre for m in proc.required_measures},
231        post_measures={m for proc in post for m in proc.required_measures},
232    )