bioimageio.core.axis

 1from __future__ import annotations
 2
 3from dataclasses import dataclass
 4from typing import Literal, Mapping, Optional, TypeVar, Union
 5
 6from typing_extensions import assert_never
 7
 8from bioimageio.spec.model import v0_5
 9
10
11def _get_axis_type(a: Literal["b", "t", "i", "c", "x", "y", "z"]):
12    if a == "b":
13        return "batch"
14    elif a == "t":
15        return "time"
16    elif a == "i":
17        return "index"
18    elif a == "c":
19        return "channel"
20    elif a in ("x", "y", "z"):
21        return "space"
22    else:
23        return "index"  # return most unspecific axis
24
25
26S = TypeVar("S", bound=str)
27
28
29AxisId = v0_5.AxisId
30
31T = TypeVar("T")
32PerAxis = Mapping[AxisId, T]
33
34BatchSize = int
35
36AxisLetter = Literal["b", "i", "t", "c", "z", "y", "x"]
37AxisLike = Union[AxisId, AxisLetter, v0_5.AnyAxis, "Axis"]
38
39
40@dataclass
41class Axis:
42    id: AxisId
43    type: Literal["batch", "channel", "index", "space", "time"]
44
45    @classmethod
46    def create(cls, axis: AxisLike) -> Axis:
47        if isinstance(axis, cls):
48            return axis
49        elif isinstance(axis, Axis):
50            return Axis(id=axis.id, type=axis.type)
51        elif isinstance(axis, str):
52            return Axis(id=AxisId(axis), type=_get_axis_type(axis))
53        elif isinstance(axis, v0_5.AxisBase):
54            return Axis(id=AxisId(axis.id), type=axis.type)
55        else:
56            assert_never(axis)
57
58
59@dataclass
60class AxisInfo(Axis):
61    maybe_singleton: bool  # TODO: replace 'maybe_singleton' with size min/max for better axis guessing
62
63    @classmethod
64    def create(cls, axis: AxisLike, maybe_singleton: Optional[bool] = None) -> AxisInfo:
65        if isinstance(axis, AxisInfo):
66            return axis
67
68        axis_base = super().create(axis)
69        if maybe_singleton is None:
70            if isinstance(axis, (Axis, str)):
71                maybe_singleton = True
72            else:
73                if axis.size is None:
74                    maybe_singleton = True
75                elif isinstance(axis.size, int):
76                    maybe_singleton = axis.size == 1
77                elif isinstance(axis.size, v0_5.SizeReference):
78                    maybe_singleton = (
79                        True  # TODO: check if singleton is ok for a `SizeReference`
80                    )
81                elif isinstance(
82                    axis.size, (v0_5.ParameterizedSize, v0_5.DataDependentSize)
83                ):
84                    try:
85                        maybe_size_one = axis.size.validate_size(
86                            1
87                        )  # TODO: refactor validate_size() to have boolean func here
88                    except ValueError:
89                        maybe_singleton = False
90                    else:
91                        maybe_singleton = maybe_size_one == 1
92                else:
93                    assert_never(axis.size)
94
95        return AxisInfo(
96            id=axis_base.id, type=axis_base.type, maybe_singleton=maybe_singleton
97        )