bioimageio.core.prediction

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