bioimageio.spec.model

2062class ModelDescr(GenericModelDescrBase, title="bioimage.io model specification"):
2063    """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
2064    These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
2065    """
2066
2067    format_version: Literal["0.5.3"] = "0.5.3"
2068    """Version of the bioimage.io model description specification used.
2069    When creating a new model always use the latest micro/patch version described here.
2070    The `format_version` is important for any consumer software to understand how to parse the fields.
2071    """
2072
2073    type: Literal["model"] = "model"
2074    """Specialized resource type 'model'"""
2075
2076    id: Optional[ModelId] = None
2077    """bioimage.io-wide unique resource identifier
2078    assigned by bioimage.io; version **un**specific."""
2079
2080    authors: NotEmpty[List[Author]]
2081    """The authors are the creators of the model RDF and the primary points of contact."""
2082
2083    documentation: Annotated[
2084        DocumentationSource,
2085        Field(
2086            examples=[
2087                "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_descriptions/models/unet2d_nuclei_broad/README.md",
2088                "README.md",
2089            ],
2090        ),
2091    ]
2092    """∈📦 URL or relative path to a markdown file with additional documentation.
2093    The recommended documentation file name is `README.md`. An `.md` suffix is mandatory.
2094    The documentation should include a '#[#] Validation' (sub)section
2095    with details on how to quantitatively validate the model on unseen data."""
2096
2097    @field_validator("documentation", mode="after")
2098    @classmethod
2099    def _validate_documentation(cls, value: DocumentationSource) -> DocumentationSource:
2100        if not validation_context_var.get().perform_io_checks:
2101            return value
2102
2103        doc_path = download(value).path
2104        doc_content = doc_path.read_text(encoding="utf-8")
2105        assert isinstance(doc_content, str)
2106        if not re.match("#.*[vV]alidation", doc_content):
2107            issue_warning(
2108                "No '# Validation' (sub)section found in {value}.",
2109                value=value,
2110                field="documentation",
2111            )
2112
2113        return value
2114
2115    inputs: NotEmpty[Sequence[InputTensorDescr]]
2116    """Describes the input tensors expected by this model."""
2117
2118    @field_validator("inputs", mode="after")
2119    @classmethod
2120    def _validate_input_axes(
2121        cls, inputs: Sequence[InputTensorDescr]
2122    ) -> Sequence[InputTensorDescr]:
2123        input_size_refs = cls._get_axes_with_independent_size(inputs)
2124
2125        for i, ipt in enumerate(inputs):
2126            valid_independent_refs: Dict[
2127                Tuple[TensorId, AxisId],
2128                Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
2129            ] = {
2130                **{
2131                    (ipt.id, a.id): (ipt, a, a.size)
2132                    for a in ipt.axes
2133                    if not isinstance(a, BatchAxis)
2134                    and isinstance(a.size, (int, ParameterizedSize))
2135                },
2136                **input_size_refs,
2137            }
2138            for a, ax in enumerate(ipt.axes):
2139                cls._validate_axis(
2140                    "inputs",
2141                    i=i,
2142                    tensor_id=ipt.id,
2143                    a=a,
2144                    axis=ax,
2145                    valid_independent_refs=valid_independent_refs,
2146                )
2147        return inputs
2148
2149    @staticmethod
2150    def _validate_axis(
2151        field_name: str,
2152        i: int,
2153        tensor_id: TensorId,
2154        a: int,
2155        axis: AnyAxis,
2156        valid_independent_refs: Dict[
2157            Tuple[TensorId, AxisId],
2158            Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
2159        ],
2160    ):
2161        if isinstance(axis, BatchAxis) or isinstance(
2162            axis.size, (int, ParameterizedSize, DataDependentSize)
2163        ):
2164            return
2165        elif not isinstance(axis.size, SizeReference):
2166            assert_never(axis.size)
2167
2168        # validate axis.size SizeReference
2169        ref = (axis.size.tensor_id, axis.size.axis_id)
2170        if ref not in valid_independent_refs:
2171            raise ValueError(
2172                "Invalid tensor axis reference at"
2173                + f" {field_name}[{i}].axes[{a}].size: {axis.size}."
2174            )
2175        if ref == (tensor_id, axis.id):
2176            raise ValueError(
2177                "Self-referencing not allowed for"
2178                + f" {field_name}[{i}].axes[{a}].size: {axis.size}"
2179            )
2180        if axis.type == "channel":
2181            if valid_independent_refs[ref][1].type != "channel":
2182                raise ValueError(
2183                    "A channel axis' size may only reference another fixed size"
2184                    + " channel axis."
2185                )
2186            if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
2187                ref_size = valid_independent_refs[ref][2]
2188                assert isinstance(ref_size, int), (
2189                    "channel axis ref (another channel axis) has to specify fixed"
2190                    + " size"
2191                )
2192                generated_channel_names = [
2193                    Identifier(axis.channel_names.format(i=i))
2194                    for i in range(1, ref_size + 1)
2195                ]
2196                axis.channel_names = generated_channel_names
2197
2198        if (ax_unit := getattr(axis, "unit", None)) != (
2199            ref_unit := getattr(valid_independent_refs[ref][1], "unit", None)
2200        ):
2201            raise ValueError(
2202                "The units of an axis and its reference axis need to match, but"
2203                + f" '{ax_unit}' != '{ref_unit}'."
2204            )
2205        ref_axis = valid_independent_refs[ref][1]
2206        if isinstance(ref_axis, BatchAxis):
2207            raise ValueError(
2208                f"Invalid reference axis '{ref_axis.id}' for {tensor_id}.{axis.id}"
2209                + " (a batch axis is not allowed as reference)."
2210            )
2211
2212        if isinstance(axis, WithHalo):
2213            min_size = axis.size.get_size(axis, ref_axis, n=0)
2214            if (min_size - 2 * axis.halo) < 1:
2215                raise ValueError(
2216                    f"axis {axis.id} with minimum size {min_size} is too small for halo"
2217                    + f" {axis.halo}."
2218                )
2219
2220            input_halo = axis.halo * axis.scale / ref_axis.scale
2221            if input_halo != int(input_halo) or input_halo % 2 == 1:
2222                raise ValueError(
2223                    f"input_halo {input_halo} (output_halo {axis.halo} *"
2224                    + f" output_scale {axis.scale} / input_scale {ref_axis.scale})"
2225                    + f" is not an even integer for {tensor_id}.{axis.id}."
2226                )
2227
2228    @model_validator(mode="after")
2229    def _validate_test_tensors(self) -> Self:
2230        if not validation_context_var.get().perform_io_checks:
2231            return self
2232
2233        test_arrays = [
2234            load_array(descr.test_tensor.download().path)
2235            for descr in chain(self.inputs, self.outputs)
2236        ]
2237        tensors = {
2238            descr.id: (descr, array)
2239            for descr, array in zip(chain(self.inputs, self.outputs), test_arrays)
2240        }
2241        validate_tensors(tensors, tensor_origin="test_tensor")
2242        return self
2243
2244    @model_validator(mode="after")
2245    def _validate_tensor_references_in_proc_kwargs(self, info: ValidationInfo) -> Self:
2246        ipt_refs = {t.id for t in self.inputs}
2247        out_refs = {t.id for t in self.outputs}
2248        for ipt in self.inputs:
2249            for p in ipt.preprocessing:
2250                ref = p.kwargs.get("reference_tensor")
2251                if ref is None:
2252                    continue
2253                if ref not in ipt_refs:
2254                    raise ValueError(
2255                        f"`reference_tensor` '{ref}' not found. Valid input tensor"
2256                        + f" references are: {ipt_refs}."
2257                    )
2258
2259        for out in self.outputs:
2260            for p in out.postprocessing:
2261                ref = p.kwargs.get("reference_tensor")
2262                if ref is None:
2263                    continue
2264
2265                if ref not in ipt_refs and ref not in out_refs:
2266                    raise ValueError(
2267                        f"`reference_tensor` '{ref}' not found. Valid tensor references"
2268                        + f" are: {ipt_refs | out_refs}."
2269                    )
2270
2271        return self
2272
2273    # TODO: use validate funcs in validate_test_tensors
2274    # def validate_inputs(self, input_tensors: Mapping[TensorId, NDArray[Any]]) -> Mapping[TensorId, NDArray[Any]]:
2275
2276    name: Annotated[
2277        Annotated[
2278            str, RestrictCharacters(string.ascii_letters + string.digits + "_- ()")
2279        ],
2280        MinLen(5),
2281        MaxLen(128),
2282        warn(MaxLen(64), "Name longer than 64 characters.", INFO),
2283    ]
2284    """A human-readable name of this model.
2285    It should be no longer than 64 characters
2286    and may only contain letter, number, underscore, minus, parentheses and spaces.
2287    We recommend to chose a name that refers to the model's task and image modality.
2288    """
2289
2290    outputs: NotEmpty[Sequence[OutputTensorDescr]]
2291    """Describes the output tensors."""
2292
2293    @field_validator("outputs", mode="after")
2294    @classmethod
2295    def _validate_tensor_ids(
2296        cls, outputs: Sequence[OutputTensorDescr], info: ValidationInfo
2297    ) -> Sequence[OutputTensorDescr]:
2298        tensor_ids = [
2299            t.id for t in info.data.get("inputs", []) + info.data.get("outputs", [])
2300        ]
2301        duplicate_tensor_ids: List[str] = []
2302        seen: Set[str] = set()
2303        for t in tensor_ids:
2304            if t in seen:
2305                duplicate_tensor_ids.append(t)
2306
2307            seen.add(t)
2308
2309        if duplicate_tensor_ids:
2310            raise ValueError(f"Duplicate tensor ids: {duplicate_tensor_ids}")
2311
2312        return outputs
2313
2314    @staticmethod
2315    def _get_axes_with_parameterized_size(
2316        io: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
2317    ):
2318        return {
2319            f"{t.id}.{a.id}": (t, a, a.size)
2320            for t in io
2321            for a in t.axes
2322            if not isinstance(a, BatchAxis) and isinstance(a.size, ParameterizedSize)
2323        }
2324
2325    @staticmethod
2326    def _get_axes_with_independent_size(
2327        io: Union[Sequence[InputTensorDescr], Sequence[OutputTensorDescr]],
2328    ):
2329        return {
2330            (t.id, a.id): (t, a, a.size)
2331            for t in io
2332            for a in t.axes
2333            if not isinstance(a, BatchAxis)
2334            and isinstance(a.size, (int, ParameterizedSize))
2335        }
2336
2337    @field_validator("outputs", mode="after")
2338    @classmethod
2339    def _validate_output_axes(
2340        cls, outputs: List[OutputTensorDescr], info: ValidationInfo
2341    ) -> List[OutputTensorDescr]:
2342        input_size_refs = cls._get_axes_with_independent_size(
2343            info.data.get("inputs", [])
2344        )
2345        output_size_refs = cls._get_axes_with_independent_size(outputs)
2346
2347        for i, out in enumerate(outputs):
2348            valid_independent_refs: Dict[
2349                Tuple[TensorId, AxisId],
2350                Tuple[TensorDescr, AnyAxis, Union[int, ParameterizedSize]],
2351            ] = {
2352                **{
2353                    (out.id, a.id): (out, a, a.size)
2354                    for a in out.axes
2355                    if not isinstance(a, BatchAxis)
2356                    and isinstance(a.size, (int, ParameterizedSize))
2357                },
2358                **input_size_refs,
2359                **output_size_refs,
2360            }
2361            for a, ax in enumerate(out.axes):
2362                cls._validate_axis(
2363                    "outputs",
2364                    i,
2365                    out.id,
2366                    a,
2367                    ax,
2368                    valid_independent_refs=valid_independent_refs,
2369                )
2370
2371        return outputs
2372
2373    packaged_by: List[Author] = Field(default_factory=list)
2374    """The persons that have packaged and uploaded this model.
2375    Only required if those persons differ from the `authors`."""
2376
2377    parent: Optional[LinkedModel] = None
2378    """The model from which this model is derived, e.g. by fine-tuning the weights."""
2379
2380    # todo: add parent self check once we have `id`
2381    # @model_validator(mode="after")
2382    # def validate_parent_is_not_self(self) -> Self:
2383    #     if self.parent is not None and self.parent == self.id:
2384    #         raise ValueError("The model may not reference itself as parent model")
2385
2386    #     return self
2387
2388    run_mode: Annotated[
2389        Optional[RunMode],
2390        warn(None, "Run mode '{value}' has limited support across consumer softwares."),
2391    ] = None
2392    """Custom run mode for this model: for more complex prediction procedures like test time
2393    data augmentation that currently cannot be expressed in the specification.
2394    No standard run modes are defined yet."""
2395
2396    timestamp: Datetime = Datetime(datetime.now())
2397    """Timestamp in [ISO 8601](#https://en.wikipedia.org/wiki/ISO_8601) format
2398    with a few restrictions listed [here](https://docs.python.org/3/library/datetime.html#datetime.datetime.fromisoformat).
2399    (In Python a datetime object is valid, too)."""
2400
2401    training_data: Annotated[
2402        Union[None, LinkedDataset, DatasetDescr, DatasetDescr02],
2403        Field(union_mode="left_to_right"),
2404    ] = None
2405    """The dataset used to train this model"""
2406
2407    weights: Annotated[WeightsDescr, WrapSerializer(package_weights)]
2408    """The weights for this model.
2409    Weights can be given for different formats, but should otherwise be equivalent.
2410    The available weight formats determine which consumers can use this model."""
2411
2412    @model_validator(mode="after")
2413    def _add_default_cover(self) -> Self:
2414        if not validation_context_var.get().perform_io_checks or self.covers:
2415            return self
2416
2417        try:
2418            generated_covers = generate_covers(
2419                [(t, load_array(t.test_tensor.download().path)) for t in self.inputs],
2420                [(t, load_array(t.test_tensor.download().path)) for t in self.outputs],
2421            )
2422        except Exception as e:
2423            issue_warning(
2424                "Failed to generate cover image(s): {e}",
2425                value=self.covers,
2426                msg_context=dict(e=e),
2427                field="covers",
2428            )
2429        else:
2430            self.covers.extend(generated_covers)
2431
2432        return self
2433
2434    def get_input_test_arrays(self) -> List[NDArray[Any]]:
2435        data = [load_array(ipt.test_tensor.download().path) for ipt in self.inputs]
2436        assert all(isinstance(d, np.ndarray) for d in data)
2437        return data
2438
2439    def get_output_test_arrays(self) -> List[NDArray[Any]]:
2440        data = [load_array(out.test_tensor.download().path) for out in self.outputs]
2441        assert all(isinstance(d, np.ndarray) for d in data)
2442        return data
2443
2444    @staticmethod
2445    def get_batch_size(tensor_sizes: Mapping[TensorId, Mapping[AxisId, int]]) -> int:
2446        batch_size = 1
2447        tensor_with_batchsize: Optional[TensorId] = None
2448        for tid in tensor_sizes:
2449            for aid, s in tensor_sizes[tid].items():
2450                if aid != BATCH_AXIS_ID or s == 1 or s == batch_size:
2451                    continue
2452
2453                if batch_size != 1:
2454                    assert tensor_with_batchsize is not None
2455                    raise ValueError(
2456                        f"batch size mismatch for tensors '{tensor_with_batchsize}' ({batch_size}) and '{tid}' ({s})"
2457                    )
2458
2459                batch_size = s
2460                tensor_with_batchsize = tid
2461
2462        return batch_size
2463
2464    def get_output_tensor_sizes(
2465        self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]
2466    ) -> Dict[TensorId, Dict[AxisId, Union[int, _DataDepSize]]]:
2467        """Returns the tensor output sizes for given **input_sizes**.
2468        Only if **input_sizes** has a valid input shape, the tensor output size is exact.
2469        Otherwise it might be larger than the actual (valid) output"""
2470        batch_size = self.get_batch_size(input_sizes)
2471        ns = self.get_ns(input_sizes)
2472
2473        tensor_sizes = self.get_tensor_sizes(ns, batch_size=batch_size)
2474        return tensor_sizes.outputs
2475
2476    def get_ns(self, input_sizes: Mapping[TensorId, Mapping[AxisId, int]]):
2477        """get parameter `n` for each parameterized axis
2478        such that the valid input size is >= the given input size"""
2479        ret: Dict[Tuple[TensorId, AxisId], ParameterizedSize_N] = {}
2480        axes = {t.id: {a.id: a for a in t.axes} for t in self.inputs}
2481        for tid in input_sizes:
2482            for aid, s in input_sizes[tid].items():
2483                size_descr = axes[tid][aid].size
2484                if isinstance(size_descr, ParameterizedSize):
2485                    ret[(tid, aid)] = size_descr.get_n(s)
2486                elif size_descr is None or isinstance(size_descr, (int, SizeReference)):
2487                    pass
2488                else:
2489                    assert_never(size_descr)
2490
2491        return ret
2492
2493    def get_tensor_sizes(
2494        self, ns: Mapping[Tuple[TensorId, AxisId], ParameterizedSize_N], batch_size: int
2495    ) -> _TensorSizes:
2496        axis_sizes = self.get_axis_sizes(ns, batch_size=batch_size)
2497        return _TensorSizes(
2498            {
2499                t: {
2500                    aa: axis_sizes.inputs[(tt, aa)]
2501                    for tt, aa in axis_sizes.inputs
2502                    if tt == t
2503                }
2504                for t in {tt for tt, _ in axis_sizes.inputs}
2505            },
2506            {
2507                t: {
2508                    aa: axis_sizes.outputs[(tt, aa)]
2509                    for tt, aa in axis_sizes.outputs
2510                    if tt == t
2511                }
2512                for t in {tt for tt, _ in axis_sizes.outputs}
2513            },
2514        )
2515
2516    def get_axis_sizes(
2517        self,
2518        ns: Mapping[Tuple[TensorId, AxisId], ParameterizedSize_N],
2519        batch_size: Optional[int] = None,
2520        *,
2521        max_input_shape: Optional[Mapping[Tuple[TensorId, AxisId], int]] = None,
2522    ) -> _AxisSizes:
2523        """Determine input and output block shape for scale factors **ns**
2524        of parameterized input sizes.
2525
2526        Args:
2527            ns: Scale factor `n` for each axis (keyed by (tensor_id, axis_id))
2528                that is parameterized as `size = min + n * step`.
2529            batch_size: The desired size of the batch dimension.
2530                If given **batch_size** overwrites any batch size present in
2531                **max_input_shape**. Default 1.
2532            max_input_shape: Limits the derived block shapes.
2533                Each axis for which the input size, parameterized by `n`, is larger
2534                than **max_input_shape** is set to the minimal value `n_min` for which
2535                this is still true.
2536                Use this for small input samples or large values of **ns**.
2537                Or simply whenever you know the full input shape.
2538
2539        Returns:
2540            Resolved axis sizes for model inputs and outputs.
2541        """
2542        max_input_shape = max_input_shape or {}
2543        if batch_size is None:
2544            for (_t_id, a_id), s in max_input_shape.items():
2545                if a_id == BATCH_AXIS_ID:
2546                    batch_size = s
2547                    break
2548            else:
2549                batch_size = 1
2550
2551        all_axes = {
2552            t.id: {a.id: a for a in t.axes} for t in chain(self.inputs, self.outputs)
2553        }
2554
2555        inputs: Dict[Tuple[TensorId, AxisId], int] = {}
2556        outputs: Dict[Tuple[TensorId, AxisId], Union[int, _DataDepSize]] = {}
2557
2558        def get_axis_size(a: Union[InputAxis, OutputAxis]):
2559            if isinstance(a, BatchAxis):
2560                if (t_descr.id, a.id) in ns:
2561                    logger.warning(
2562                        "Ignoring unexpected size increment factor (n) for batch axis"
2563                        + " of tensor '{}'.",
2564                        t_descr.id,
2565                    )
2566                return batch_size
2567            elif isinstance(a.size, int):
2568                if (t_descr.id, a.id) in ns:
2569                    logger.warning(
2570                        "Ignoring unexpected size increment factor (n) for fixed size"
2571                        + " axis '{}' of tensor '{}'.",
2572                        a.id,
2573                        t_descr.id,
2574                    )
2575                return a.size
2576            elif isinstance(a.size, ParameterizedSize):
2577                if (t_descr.id, a.id) not in ns:
2578                    raise ValueError(
2579                        "Size increment factor (n) missing for parametrized axis"
2580                        + f" '{a.id}' of tensor '{t_descr.id}'."
2581                    )
2582                n = ns[(t_descr.id, a.id)]
2583                s_max = max_input_shape.get((t_descr.id, a.id))
2584                if s_max is not None:
2585                    n = min(n, a.size.get_n(s_max))
2586
2587                return a.size.get_size(n)
2588
2589            elif isinstance(a.size, SizeReference):
2590                if (t_descr.id, a.id) in ns:
2591                    logger.warning(
2592                        "Ignoring unexpected size increment factor (n) for axis '{}'"
2593                        + " of tensor '{}' with size reference.",
2594                        a.id,
2595                        t_descr.id,
2596                    )
2597                assert not isinstance(a, BatchAxis)
2598                ref_axis = all_axes[a.size.tensor_id][a.size.axis_id]
2599                assert not isinstance(ref_axis, BatchAxis)
2600                ref_key = (a.size.tensor_id, a.size.axis_id)
2601                ref_size = inputs.get(ref_key, outputs.get(ref_key))
2602                assert ref_size is not None, ref_key
2603                assert not isinstance(ref_size, _DataDepSize), ref_key
2604                return a.size.get_size(
2605                    axis=a,
2606                    ref_axis=ref_axis,
2607                    ref_size=ref_size,
2608                )
2609            elif isinstance(a.size, DataDependentSize):
2610                if (t_descr.id, a.id) in ns:
2611                    logger.warning(
2612                        "Ignoring unexpected increment factor (n) for data dependent"
2613                        + " size axis '{}' of tensor '{}'.",
2614                        a.id,
2615                        t_descr.id,
2616                    )
2617                return _DataDepSize(a.size.min, a.size.max)
2618            else:
2619                assert_never(a.size)
2620
2621        # first resolve all , but the `SizeReference` input sizes
2622        for t_descr in self.inputs:
2623            for a in t_descr.axes:
2624                if not isinstance(a.size, SizeReference):
2625                    s = get_axis_size(a)
2626                    assert not isinstance(s, _DataDepSize)
2627                    inputs[t_descr.id, a.id] = s
2628
2629        # resolve all other input axis sizes
2630        for t_descr in self.inputs:
2631            for a in t_descr.axes:
2632                if isinstance(a.size, SizeReference):
2633                    s = get_axis_size(a)
2634                    assert not isinstance(s, _DataDepSize)
2635                    inputs[t_descr.id, a.id] = s
2636
2637        # resolve all output axis sizes
2638        for t_descr in self.outputs:
2639            for a in t_descr.axes:
2640                assert not isinstance(a.size, ParameterizedSize)
2641                s = get_axis_size(a)
2642                outputs[t_descr.id, a.id] = s
2643
2644        return _AxisSizes(inputs=inputs, outputs=outputs)
2645
2646    @model_validator(mode="before")
2647    @classmethod
2648    def _convert(cls, data: Dict[str, Any]) -> Dict[str, Any]:
2649        if (
2650            data.get("type") == "model"
2651            and isinstance(fv := data.get("format_version"), str)
2652            and fv.count(".") == 2
2653        ):
2654            fv_parts = fv.split(".")
2655            if any(not p.isdigit() for p in fv_parts):
2656                return data
2657
2658            fv_tuple = tuple(map(int, fv_parts))
2659
2660            assert cls.implemented_format_version_tuple[0:2] == (0, 5)
2661            if fv_tuple[:2] in ((0, 3), (0, 4)):
2662                m04 = _ModelDescr_v0_4.load(data)
2663                if not isinstance(m04, InvalidDescr):
2664                    return _model_conv.convert_as_dict(m04)
2665            elif fv_tuple[:2] == (0, 5):
2666                # bump patch version
2667                data["format_version"] = cls.implemented_format_version
2668
2669        return data