bioimageio.core.prediction

  1import collections.abc
  2from pathlib import Path
  3from typing import (
  4    Hashable,
  5    Iterable,
  6    Iterator,
  7    Mapping,
  8    Optional,
  9    Tuple,
 10    Union,
 11)
 12
 13from loguru import logger
 14from tqdm import tqdm
 15
 16from bioimageio.spec import load_description
 17from bioimageio.spec.common import PermissiveFileSource
 18from bioimageio.spec.model import v0_4, v0_5
 19
 20from ._prediction_pipeline import PredictionPipeline, create_prediction_pipeline
 21from .axis import AxisId
 22from .common import BlocksizeParameter, MemberId, PerMember
 23from .digest_spec import TensorSource, create_sample_for_model, get_member_id
 24from .io import save_sample
 25from .sample import Sample
 26
 27
 28def predict(
 29    *,
 30    model: Union[
 31        PermissiveFileSource, v0_4.ModelDescr, v0_5.ModelDescr, PredictionPipeline
 32    ],
 33    inputs: Union[Sample, PerMember[TensorSource], TensorSource],
 34    sample_id: Hashable = "sample",
 35    blocksize_parameter: Optional[BlocksizeParameter] = None,
 36    input_block_shape: Optional[Mapping[MemberId, Mapping[AxisId, int]]] = None,
 37    skip_preprocessing: bool = False,
 38    skip_postprocessing: bool = False,
 39    save_output_path: Optional[Union[Path, str]] = None,
 40) -> Sample:
 41    """Run prediction for a single set of input(s) with a bioimage.io model
 42
 43    Args:
 44        model: Model to predict with.
 45            May be given as RDF source, model description or prediction pipeline.
 46        inputs: the input sample or the named input(s) for this model as a dictionary
 47        sample_id: the sample id.
 48            The **sample_id** is used to format **save_output_path**
 49            and to distinguish sample specific log messages.
 50        blocksize_parameter: (optional) Tile the input into blocks parametrized by
 51            **blocksize_parameter** according to any parametrized axis sizes defined
 52            by the **model**.
 53            See `bioimageio.spec.model.v0_5.ParameterizedSize` for details.
 54            Note: For a predetermined, fixed block shape use **input_block_shape**.
 55        input_block_shape: (optional) Tile the input sample tensors into blocks.
 56            Note: Use **blocksize_parameter** for a parameterized block shape to
 57                run prediction independent of the exact block shape.
 58        skip_preprocessing: Flag to skip the model's preprocessing.
 59        skip_postprocessing: Flag to skip the model's postprocessing.
 60        save_output_path: A path with to save the output to. M
 61            Must contain:
 62            - `{output_id}` (or `{member_id}`) if the model has multiple output tensors
 63            May contain:
 64            - `{sample_id}` to avoid overwriting recurrent calls
 65    """
 66    if isinstance(model, PredictionPipeline):
 67        pp = model
 68        model = pp.model_description
 69    else:
 70        if not isinstance(model, (v0_4.ModelDescr, v0_5.ModelDescr)):
 71            loaded = load_description(model)
 72            if not isinstance(loaded, (v0_4.ModelDescr, v0_5.ModelDescr)):
 73                raise ValueError(f"expected model description, but got {loaded}")
 74            model = loaded
 75
 76        pp = create_prediction_pipeline(model)
 77
 78    if save_output_path is not None:
 79        if (
 80            "{output_id}" not in str(save_output_path)
 81            and "{member_id}" not in str(save_output_path)
 82            and len(model.outputs) > 1
 83        ):
 84            raise ValueError(
 85                f"Missing `{{output_id}}` in save_output_path={save_output_path} to "
 86                + "distinguish model outputs "
 87                + str([get_member_id(d) for d in model.outputs])
 88            )
 89
 90    if isinstance(inputs, Sample):
 91        sample = inputs
 92    else:
 93        sample = create_sample_for_model(
 94            pp.model_description, inputs=inputs, sample_id=sample_id
 95        )
 96
 97    if input_block_shape is not None:
 98        if blocksize_parameter is not None:
 99            logger.warning(
100                "ignoring blocksize_parameter={} in favor of input_block_shape={}",
101                blocksize_parameter,
102                input_block_shape,
103            )
104
105        output = pp.predict_sample_with_fixed_blocking(
106            sample,
107            input_block_shape=input_block_shape,
108            skip_preprocessing=skip_preprocessing,
109            skip_postprocessing=skip_postprocessing,
110        )
111    elif blocksize_parameter is not None:
112        output = pp.predict_sample_with_blocking(
113            sample,
114            skip_preprocessing=skip_preprocessing,
115            skip_postprocessing=skip_postprocessing,
116            ns=blocksize_parameter,
117        )
118    else:
119        output = pp.predict_sample_without_blocking(
120            sample,
121            skip_preprocessing=skip_preprocessing,
122            skip_postprocessing=skip_postprocessing,
123        )
124    if save_output_path:
125        save_sample(save_output_path, output)
126
127    return output
128
129
130def predict_many(
131    *,
132    model: Union[
133        PermissiveFileSource, v0_4.ModelDescr, v0_5.ModelDescr, PredictionPipeline
134    ],
135    inputs: Union[Iterable[PerMember[TensorSource]], Iterable[TensorSource]],
136    sample_id: str = "sample{i:03}",
137    blocksize_parameter: Optional[
138        Union[
139            v0_5.ParameterizedSize_N,
140            Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N],
141        ]
142    ] = None,
143    skip_preprocessing: bool = False,
144    skip_postprocessing: bool = False,
145    save_output_path: Optional[Union[Path, str]] = None,
146) -> Iterator[Sample]:
147    """Run prediction for a multiple sets of inputs with a bioimage.io model
148
149    Args:
150        model: Model to predict with.
151            May be given as RDF source, model description or prediction pipeline.
152        inputs: An iterable of the named input(s) for this model as a dictionary.
153        sample_id: The sample id.
154            note: `{i}` will be formatted as the i-th sample.
155            If `{i}` (or `{i:`) is not present and `inputs` is not an iterable `{i:03}`
156            is appended.
157        blocksize_parameter: (optional) Tile the input into blocks parametrized by
158            blocksize according to any parametrized axis sizes defined in the model RDF.
159        skip_preprocessing: Flag to skip the model's preprocessing.
160        skip_postprocessing: Flag to skip the model's postprocessing.
161        save_output_path: A path to save the output to.
162            Must contain:
163            - `{sample_id}` to differentiate predicted samples
164            - `{output_id}` (or `{member_id}`) if the model has multiple outputs
165    """
166    if save_output_path is not None and "{sample_id}" not in str(save_output_path):
167        raise ValueError(
168            f"Missing `{{sample_id}}` in save_output_path={save_output_path}"
169            + " to differentiate predicted samples."
170        )
171
172    if isinstance(model, PredictionPipeline):
173        pp = model
174    else:
175        if not isinstance(model, (v0_4.ModelDescr, v0_5.ModelDescr)):
176            loaded = load_description(model)
177            if not isinstance(loaded, (v0_4.ModelDescr, v0_5.ModelDescr)):
178                raise ValueError(f"expected model description, but got {loaded}")
179            model = loaded
180
181        pp = create_prediction_pipeline(model)
182
183    if not isinstance(inputs, collections.abc.Mapping):
184        if "{i}" not in sample_id and "{i:" not in sample_id:
185            sample_id += "{i:03}"
186
187        total = len(inputs) if isinstance(inputs, collections.abc.Sized) else None
188
189        for i, ipts in tqdm(enumerate(inputs), total=total):
190            yield predict(
191                model=pp,
192                inputs=ipts,
193                sample_id=sample_id.format(i=i),
194                blocksize_parameter=blocksize_parameter,
195                skip_preprocessing=skip_preprocessing,
196                skip_postprocessing=skip_postprocessing,
197                save_output_path=save_output_path,
198            )