Coverage for bioimageio/spec/_package.py: 81%
88 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-27 09:20 +0000
1import collections.abc
2import shutil
3from io import BytesIO
4from pathlib import Path
5from tempfile import NamedTemporaryFile, mkdtemp
6from typing import IO, Dict, Literal, Optional, Sequence, Union
7from zipfile import ZIP_DEFLATED
9from loguru import logger
10from pydantic import DirectoryPath, FilePath, NewPath
12from ._description import InvalidDescr, ResourceDescr, build_description
13from ._internal.common_nodes import ResourceDescrBase
14from ._internal.io import (
15 BioimageioYamlContent,
16 BioimageioYamlSource,
17 FileDescr,
18 RelativeFilePath,
19 ensure_is_valid_bioimageio_yaml_name,
20)
21from ._internal.io_basics import (
22 BIOIMAGEIO_YAML,
23 AbsoluteFilePath,
24 BytesReader,
25 FileName,
26 ZipPath,
27)
28from ._internal.io_utils import open_bioimageio_yaml, write_yaml, write_zip
29from ._internal.packaging_context import PackagingContext
30from ._internal.url import HttpUrl
31from ._internal.utils import get_os_friendly_file_name
32from ._internal.validation_context import get_validation_context
33from ._internal.warning_levels import ERROR
34from ._io import load_description
35from .model.v0_4 import WeightsFormat
38# TODO: deprecate in favor of get_package_content
39def get_resource_package_content(
40 rd: ResourceDescr,
41 /,
42 *,
43 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
44 weights_priority_order: Optional[Sequence[WeightsFormat]] = None, # model only
45) -> Dict[FileName, Union[HttpUrl, AbsoluteFilePath, BioimageioYamlContent, ZipPath]]:
46 ret: Dict[
47 FileName, Union[HttpUrl, AbsoluteFilePath, BioimageioYamlContent, ZipPath]
48 ] = {}
49 for k, v in get_package_content(
50 rd,
51 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
52 weights_priority_order=weights_priority_order,
53 ).items():
54 if isinstance(v, FileDescr):
55 if isinstance(v.source, (Path, RelativeFilePath)):
56 ret[k] = v.source.absolute()
57 else:
58 ret[k] = v.source
60 else:
61 ret[k] = v
63 return ret
66def get_package_content(
67 rd: ResourceDescr,
68 /,
69 *,
70 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
71 weights_priority_order: Optional[Sequence[WeightsFormat]] = None, # model only
72) -> Dict[FileName, Union[FileDescr, BioimageioYamlContent]]:
73 """
74 Args:
75 rd: resource description
76 bioimageio_yaml_file_name: RDF file name
77 # for model resources only:
78 weights_priority_order: If given, only the first weights format present in the model is included.
79 If none of the prioritized weights formats is found a ValueError is raised.
80 """
81 os_friendly_name = get_os_friendly_file_name(rd.name)
82 bioimageio_yaml_file_name = bioimageio_yaml_file_name.format(
83 name=os_friendly_name, type=rd.type
84 )
86 bioimageio_yaml_file_name = ensure_is_valid_bioimageio_yaml_name(
87 bioimageio_yaml_file_name
88 )
89 content: Dict[FileName, FileDescr] = {}
90 with PackagingContext(
91 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
92 file_sources=content,
93 weights_priority_order=weights_priority_order,
94 ):
95 rdf_content: BioimageioYamlContent = rd.model_dump(
96 mode="json", exclude_unset=True
97 )
99 _ = rdf_content.pop("rdf_source", None)
101 return {**content, bioimageio_yaml_file_name: rdf_content}
104def _prepare_resource_package(
105 source: Union[BioimageioYamlSource, ResourceDescr],
106 /,
107 *,
108 weights_priority_order: Optional[Sequence[WeightsFormat]] = None,
109) -> Dict[FileName, Union[BioimageioYamlContent, BytesReader]]:
110 """Prepare to package a resource description; downloads all required files.
112 Args:
113 source: A bioimage.io resource description (as file, raw YAML content or description class)
114 context: validation context
115 weights_priority_order: If given only the first weights format present in the model is included.
116 If none of the prioritized weights formats is found all are included.
117 """
118 context = get_validation_context()
119 bioimageio_yaml_file_name = context.file_name
120 if isinstance(source, ResourceDescrBase):
121 descr = source
122 elif isinstance(source, collections.abc.Mapping):
123 descr = build_description(source)
124 else:
125 opened = open_bioimageio_yaml(source)
126 bioimageio_yaml_file_name = opened.original_file_name
127 context = context.replace(
128 root=opened.original_root, file_name=opened.original_file_name
129 )
130 with context:
131 descr = build_description(opened.content)
133 if isinstance(descr, InvalidDescr):
134 raise ValueError(f"{source} is invalid: {descr.validation_summary}")
136 with context:
137 package_content = get_package_content(
138 descr,
139 bioimageio_yaml_file_name=bioimageio_yaml_file_name or BIOIMAGEIO_YAML,
140 weights_priority_order=weights_priority_order,
141 )
143 return {
144 k: v if isinstance(v, collections.abc.Mapping) else v.get_reader()
145 for k, v in package_content.items()
146 }
149def save_bioimageio_package_as_folder(
150 source: Union[BioimageioYamlSource, ResourceDescr],
151 /,
152 *,
153 output_path: Union[NewPath, DirectoryPath, None] = None,
154 weights_priority_order: Optional[ # model only
155 Sequence[
156 Literal[
157 "keras_hdf5",
158 "onnx",
159 "pytorch_state_dict",
160 "tensorflow_js",
161 "tensorflow_saved_model_bundle",
162 "torchscript",
163 ]
164 ]
165 ] = None,
166) -> DirectoryPath:
167 """Write the content of a bioimage.io resource package to a folder.
169 Args:
170 source: bioimageio resource description
171 output_path: file path to write package to
172 weights_priority_order: If given only the first weights format present in the model is included.
173 If none of the prioritized weights formats is found all are included.
175 Returns:
176 directory path to bioimageio package folder
177 """
178 package_content = _prepare_resource_package(
179 source,
180 weights_priority_order=weights_priority_order,
181 )
182 if output_path is None:
183 output_path = Path(mkdtemp())
184 else:
185 output_path = Path(output_path)
187 output_path.mkdir(exist_ok=True, parents=True)
188 for name, src in package_content.items():
189 if isinstance(src, collections.abc.Mapping):
190 write_yaml(src, output_path / name)
191 else:
192 with (output_path / name).open("wb") as dest:
193 _ = shutil.copyfileobj(src, dest)
195 return output_path
198def save_bioimageio_package(
199 source: Union[BioimageioYamlSource, ResourceDescr],
200 /,
201 *,
202 compression: int = ZIP_DEFLATED,
203 compression_level: int = 1,
204 output_path: Union[NewPath, FilePath, None] = None,
205 weights_priority_order: Optional[ # model only
206 Sequence[
207 Literal[
208 "keras_hdf5",
209 "onnx",
210 "pytorch_state_dict",
211 "tensorflow_js",
212 "tensorflow_saved_model_bundle",
213 "torchscript",
214 ]
215 ]
216 ] = None,
217 allow_invalid: bool = False,
218) -> FilePath:
219 """Package a bioimageio resource as a zip file.
221 Args:
222 rd: bioimageio resource description
223 compression: The numeric constant of compression method.
224 compression_level: Compression level to use when writing files to the archive.
225 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
226 output_path: file path to write package to
227 weights_priority_order: If given only the first weights format present in the model is included.
228 If none of the prioritized weights formats is found all are included.
230 Returns:
231 path to zipped bioimageio package
232 """
233 package_content = _prepare_resource_package(
234 source,
235 weights_priority_order=weights_priority_order,
236 )
237 if output_path is None:
238 output_path = Path(
239 NamedTemporaryFile(suffix=".bioimageio.zip", delete=False).name
240 )
241 else:
242 output_path = Path(output_path)
244 write_zip(
245 output_path,
246 package_content,
247 compression=compression,
248 compression_level=compression_level,
249 )
250 with get_validation_context().replace(warning_level=ERROR):
251 if isinstance((exported := load_description(output_path)), InvalidDescr):
252 exported.validation_summary.display()
253 msg = f"Exported package at '{output_path}' is invalid."
254 if allow_invalid:
255 logger.error(msg)
256 else:
257 raise ValueError(msg)
259 return output_path
262def save_bioimageio_package_to_stream(
263 source: Union[BioimageioYamlSource, ResourceDescr],
264 /,
265 *,
266 compression: int = ZIP_DEFLATED,
267 compression_level: int = 1,
268 output_stream: Union[IO[bytes], None] = None,
269 weights_priority_order: Optional[ # model only
270 Sequence[
271 Literal[
272 "keras_hdf5",
273 "onnx",
274 "pytorch_state_dict",
275 "tensorflow_js",
276 "tensorflow_saved_model_bundle",
277 "torchscript",
278 ]
279 ]
280 ] = None,
281) -> IO[bytes]:
282 """Package a bioimageio resource into a stream.
284 Args:
285 rd: bioimageio resource description
286 compression: The numeric constant of compression method.
287 compression_level: Compression level to use when writing files to the archive.
288 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
289 output_stream: stream to write package to
290 weights_priority_order: If given only the first weights format present in the model is included.
291 If none of the prioritized weights formats is found all are included.
293 Note: this function bypasses safety checks and does not load/validate the model after writing.
295 Returns:
296 stream of zipped bioimageio package
297 """
298 if output_stream is None:
299 output_stream = BytesIO()
301 package_content = _prepare_resource_package(
302 source,
303 weights_priority_order=weights_priority_order,
304 )
306 write_zip(
307 output_stream,
308 package_content,
309 compression=compression,
310 compression_level=compression_level,
311 )
313 return output_stream