bioimageio.core.io

  1import collections.abc
  2import warnings
  3import zipfile
  4from pathlib import Path, PurePosixPath
  5from shutil import copyfileobj
  6from typing import (
  7    Any,
  8    Mapping,
  9    Optional,
 10    Sequence,
 11    Tuple,
 12    TypeVar,
 13    Union,
 14)
 15
 16import h5py  # pyright: ignore[reportMissingTypeStubs]
 17from imageio.v3 import imread, imwrite  # type: ignore
 18from loguru import logger
 19from numpy.typing import NDArray
 20from pydantic import BaseModel, ConfigDict, TypeAdapter
 21from typing_extensions import assert_never
 22
 23from bioimageio.spec._internal.io import get_reader, interprete_file_source
 24from bioimageio.spec._internal.type_guards import is_ndarray
 25from bioimageio.spec.common import (
 26    BytesReader,
 27    FileSource,
 28    HttpUrl,
 29    PermissiveFileSource,
 30    RelativeFilePath,
 31    ZipPath,
 32)
 33from bioimageio.spec.utils import download, load_array, save_array
 34
 35from .axis import AxisLike
 36from .common import PerMember
 37from .sample import Sample
 38from .stat_measures import DatasetMeasure, MeasureValue
 39from .tensor import Tensor
 40
 41DEFAULT_H5_DATASET_PATH = "data"
 42
 43
 44SUFFIXES_WITH_DATAPATH = (".h5", ".hdf", ".hdf5")
 45
 46
 47def load_image(
 48    source: Union[ZipPath, PermissiveFileSource], is_volume: Optional[bool] = None
 49) -> NDArray[Any]:
 50    """load a single image as numpy array
 51
 52    Args:
 53        source: image source
 54        is_volume: deprecated
 55    """
 56    if is_volume is not None:
 57        warnings.warn("**is_volume** is deprecated and will be removed soon.")
 58
 59    if isinstance(source, ZipPath):
 60        parsed_source = source
 61    else:
 62        parsed_source = interprete_file_source(source)
 63
 64    if isinstance(parsed_source, RelativeFilePath):
 65        src = parsed_source.absolute()
 66    else:
 67        src = parsed_source
 68
 69    if isinstance(src, Path):
 70        file_source, suffix, subpath = _split_dataset_path(src)
 71    elif isinstance(src, HttpUrl):
 72        file_source, suffix, subpath = _split_dataset_path(src)
 73    elif isinstance(src, ZipPath):
 74        file_source, suffix, subpath = _split_dataset_path(src)
 75    else:
 76        assert_never(src)
 77
 78    if suffix == ".npy":
 79        if subpath is not None:
 80            logger.warning(
 81                "Unexpected subpath {} for .npy source {}", subpath, file_source
 82            )
 83
 84        image = load_array(file_source)
 85    elif suffix in SUFFIXES_WITH_DATAPATH:
 86        if subpath is None:
 87            dataset_path = DEFAULT_H5_DATASET_PATH
 88        else:
 89            dataset_path = str(subpath)
 90
 91        reader = download(file_source)
 92
 93        with h5py.File(reader, "r") as f:
 94            h5_dataset = f.get(  # pyright: ignore[reportUnknownVariableType]
 95                dataset_path
 96            )
 97            if not isinstance(h5_dataset, h5py.Dataset):
 98                raise ValueError(
 99                    f"{file_source} did not load as {h5py.Dataset}, but has type "
100                    + str(
101                        type(h5_dataset)  # pyright: ignore[reportUnknownArgumentType]
102                    )
103                )
104            image: NDArray[Any]
105            image = h5_dataset[:]  # pyright: ignore[reportUnknownVariableType]
106    else:
107        reader = download(file_source)
108        image = imread(  # pyright: ignore[reportUnknownVariableType]
109            reader.read(), extension=suffix
110        )
111
112    assert is_ndarray(image)
113    return image
114
115
116def load_tensor(
117    path: Union[ZipPath, Path, str], axes: Optional[Sequence[AxisLike]] = None
118) -> Tensor:
119    # TODO: load axis meta data
120    array = load_image(path)
121
122    return Tensor.from_numpy(array, dims=axes)
123
124
125_SourceT = TypeVar("_SourceT", Path, HttpUrl, ZipPath)
126
127Suffix = str
128
129
130def _split_dataset_path(
131    source: _SourceT,
132) -> Tuple[_SourceT, Suffix, Optional[PurePosixPath]]:
133    """Split off subpath (e.g. internal  h5 dataset path)
134    from a file path following a file extension.
135
136    Examples:
137        >>> _split_dataset_path(Path("my_file.h5/dataset"))
138        (...Path('my_file.h5'), '.h5', PurePosixPath('dataset'))
139
140        >>> _split_dataset_path(Path("my_plain_file"))
141        (...Path('my_plain_file'), '', None)
142
143    """
144    if isinstance(source, RelativeFilePath):
145        src = source.absolute()
146    else:
147        src = source
148
149    del source
150
151    def separate_pure_path(path: PurePosixPath):
152        for p in path.parents:
153            if p.suffix in SUFFIXES_WITH_DATAPATH:
154                return p, p.suffix, PurePosixPath(path.relative_to(p))
155
156        return path, path.suffix, None
157
158    if isinstance(src, HttpUrl):
159        file_path, suffix, data_path = separate_pure_path(PurePosixPath(src.path or ""))
160
161        if data_path is None:
162            return src, suffix, None
163
164        return (
165            HttpUrl(str(file_path).replace(f"/{data_path}", "")),
166            suffix,
167            data_path,
168        )
169
170    if isinstance(src, ZipPath):
171        file_path, suffix, data_path = separate_pure_path(PurePosixPath(str(src)))
172
173        if data_path is None:
174            return src, suffix, None
175
176        return (
177            ZipPath(str(file_path).replace(f"/{data_path}", "")),
178            suffix,
179            data_path,
180        )
181
182    file_path, suffix, data_path = separate_pure_path(PurePosixPath(src))
183    return Path(file_path), suffix, data_path
184
185
186def save_tensor(path: Union[Path, str], tensor: Tensor) -> None:
187    # TODO: save axis meta data
188
189    data: NDArray[Any] = (  # pyright: ignore[reportUnknownVariableType]
190        tensor.data.to_numpy()
191    )
192    assert is_ndarray(data)
193    file_path, suffix, subpath = _split_dataset_path(Path(path))
194    if not suffix:
195        raise ValueError(f"No suffix (needed to decide file format) found in {path}")
196
197    file_path.parent.mkdir(exist_ok=True, parents=True)
198    if file_path.suffix == ".npy":
199        if subpath is not None:
200            raise ValueError(f"Unexpected subpath {subpath} found in .npy path {path}")
201        save_array(file_path, data)
202    elif suffix in (".h5", ".hdf", ".hdf5"):
203        if subpath is None:
204            dataset_path = DEFAULT_H5_DATASET_PATH
205        else:
206            dataset_path = str(subpath)
207
208        with h5py.File(file_path, "a") as f:
209            if dataset_path in f:
210                del f[dataset_path]
211
212            _ = f.create_dataset(dataset_path, data=data, chunks=True)
213    else:
214        # if singleton_axes := [a for a, s in tensor.tagged_shape.items() if s == 1]:
215        #     tensor = tensor[{a: 0 for a in singleton_axes}]
216        #     singleton_axes_msg = f"(without singleton axes {singleton_axes}) "
217        # else:
218        singleton_axes_msg = ""
219
220        logger.debug(
221            "writing tensor {} {}to {}",
222            dict(tensor.tagged_shape),
223            singleton_axes_msg,
224            path,
225        )
226        imwrite(path, data)
227
228
229def save_sample(
230    path: Union[Path, str, PerMember[Union[Path, str]]], sample: Sample
231) -> None:
232    """Save a **sample** to a **path** pattern
233    or all sample members in the **path** mapping.
234
235    If **path** is a pathlib.Path or a string and the **sample** has multiple members,
236    **path** it must contain `{member_id}` (or `{input_id}` or `{output_id}`).
237
238    (Each) **path** may contain `{sample_id}` to be formatted with the **sample** object.
239    """
240    if not isinstance(path, collections.abc.Mapping):
241        if len(sample.members) < 2 or any(
242            m in str(path) for m in ("{member_id}", "{input_id}", "{output_id}")
243        ):
244            path = {m: path for m in sample.members}
245        else:
246            raise ValueError(
247                f"path {path} must contain '{{member_id}}' for sample with multiple members {list(sample.members)}."
248            )
249
250    for m, p in path.items():
251        t = sample.members[m]
252        p_formatted = Path(
253            str(p).format(sample_id=sample.id, member_id=m, input_id=m, output_id=m)
254        )
255        save_tensor(p_formatted, t)
256
257
258class _SerializedDatasetStatsEntry(
259    BaseModel, frozen=True, arbitrary_types_allowed=True
260):
261    measure: DatasetMeasure
262    value: MeasureValue
263
264
265_stat_adapter = TypeAdapter(
266    Sequence[_SerializedDatasetStatsEntry],
267    config=ConfigDict(arbitrary_types_allowed=True),
268)
269
270
271def save_dataset_stat(stat: Mapping[DatasetMeasure, MeasureValue], path: Path):
272    serializable = [
273        _SerializedDatasetStatsEntry(measure=k, value=v) for k, v in stat.items()
274    ]
275    _ = path.write_bytes(_stat_adapter.dump_json(serializable))
276
277
278def load_dataset_stat(path: Path):
279    seq = _stat_adapter.validate_json(path.read_bytes())
280    return {e.measure: e.value for e in seq}
281
282
283def ensure_unzipped(
284    source: Union[PermissiveFileSource, ZipPath, BytesReader], folder: Path
285):
286    """unzip a (downloaded) **source** to a file in **folder** if source is a zip archive
287    otherwise copy **source** to a file in **folder**."""
288    if isinstance(source, BytesReader):
289        weights_reader = source
290    else:
291        weights_reader = get_reader(source)
292
293    out_path = folder / (
294        weights_reader.original_file_name or f"file{weights_reader.suffix}"
295    )
296
297    if zipfile.is_zipfile(weights_reader):
298        out_path = out_path.with_name(out_path.name + ".unzipped")
299        out_path.parent.mkdir(exist_ok=True, parents=True)
300        # source itself is a zipfile
301        with zipfile.ZipFile(weights_reader, "r") as f:
302            f.extractall(out_path)
303
304    else:
305        out_path.parent.mkdir(exist_ok=True, parents=True)
306        with out_path.open("wb") as f:
307            copyfileobj(weights_reader, f)
308
309    return out_path
310
311
312def get_suffix(source: Union[ZipPath, FileSource]) -> str:
313    if isinstance(source, Path):
314        return source.suffix
315    elif isinstance(source, ZipPath):
316        return source.suffix
317    if isinstance(source, RelativeFilePath):
318        return source.path.suffix
319    elif isinstance(source, ZipPath):
320        return source.suffix
321    elif isinstance(source, HttpUrl):
322        if source.path is None:
323            return ""
324        else:
325            return PurePosixPath(source.path).suffix
326    else:
327        assert_never(source)