Coverage for bioimageio/spec/_package.py: 81%
93 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-11 07:34 +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 elif (
192 isinstance(src.original_root, Path)
193 and src.original_root / src.original_file_name
194 == (output_path / name).resolve()
195 ):
196 logger.debug(
197 f"Not copying {src.original_root / src.original_file_name} to itself."
198 )
199 else:
200 if isinstance(src.original_root, Path):
201 logger.debug(
202 f"Copying from path {src.original_root / src.original_file_name} to {output_path / name}."
203 )
204 else:
205 logger.debug(
206 f"Copying {src.original_root}/{src.original_file_name} to {output_path / name}."
207 )
208 with (output_path / name).open("wb") as dest:
209 _ = shutil.copyfileobj(src, dest)
211 return output_path
214def save_bioimageio_package(
215 source: Union[BioimageioYamlSource, ResourceDescr],
216 /,
217 *,
218 compression: int = ZIP_DEFLATED,
219 compression_level: int = 1,
220 output_path: Union[NewPath, FilePath, None] = None,
221 weights_priority_order: Optional[ # model only
222 Sequence[
223 Literal[
224 "keras_hdf5",
225 "onnx",
226 "pytorch_state_dict",
227 "tensorflow_js",
228 "tensorflow_saved_model_bundle",
229 "torchscript",
230 ]
231 ]
232 ] = None,
233 allow_invalid: bool = False,
234) -> FilePath:
235 """Package a bioimageio resource as a zip file.
237 Args:
238 rd: bioimageio resource description
239 compression: The numeric constant of compression method.
240 compression_level: Compression level to use when writing files to the archive.
241 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
242 output_path: file path to write package to
243 weights_priority_order: If given only the first weights format present in the model is included.
244 If none of the prioritized weights formats is found all are included.
246 Returns:
247 path to zipped bioimageio package
248 """
249 package_content = _prepare_resource_package(
250 source,
251 weights_priority_order=weights_priority_order,
252 )
253 if output_path is None:
254 output_path = Path(
255 NamedTemporaryFile(suffix=".bioimageio.zip", delete=False).name
256 )
257 else:
258 output_path = Path(output_path)
260 write_zip(
261 output_path,
262 package_content,
263 compression=compression,
264 compression_level=compression_level,
265 )
266 with get_validation_context().replace(warning_level=ERROR):
267 if isinstance((exported := load_description(output_path)), InvalidDescr):
268 exported.validation_summary.display()
269 msg = f"Exported package at '{output_path}' is invalid."
270 if allow_invalid:
271 logger.error(msg)
272 else:
273 raise ValueError(msg)
275 return output_path
278def save_bioimageio_package_to_stream(
279 source: Union[BioimageioYamlSource, ResourceDescr],
280 /,
281 *,
282 compression: int = ZIP_DEFLATED,
283 compression_level: int = 1,
284 output_stream: Union[IO[bytes], None] = None,
285 weights_priority_order: Optional[ # model only
286 Sequence[
287 Literal[
288 "keras_hdf5",
289 "onnx",
290 "pytorch_state_dict",
291 "tensorflow_js",
292 "tensorflow_saved_model_bundle",
293 "torchscript",
294 ]
295 ]
296 ] = None,
297) -> IO[bytes]:
298 """Package a bioimageio resource into a stream.
300 Args:
301 rd: bioimageio resource description
302 compression: The numeric constant of compression method.
303 compression_level: Compression level to use when writing files to the archive.
304 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
305 output_stream: stream to write package to
306 weights_priority_order: If given only the first weights format present in the model is included.
307 If none of the prioritized weights formats is found all are included.
309 Note: this function bypasses safety checks and does not load/validate the model after writing.
311 Returns:
312 stream of zipped bioimageio package
313 """
314 if output_stream is None:
315 output_stream = BytesIO()
317 package_content = _prepare_resource_package(
318 source,
319 weights_priority_order=weights_priority_order,
320 )
322 write_zip(
323 output_stream,
324 package_content,
325 compression=compression,
326 compression_level=compression_level,
327 )
329 return output_stream