Coverage for bioimageio/spec/_package.py: 88%
88 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-02 14:21 +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 pydantic import DirectoryPath, FilePath, NewPath
11from ._description import InvalidDescr, ResourceDescr, build_description
12from ._internal.common_nodes import ResourceDescrBase
13from ._internal.io import (
14 BioimageioYamlContent,
15 BioimageioYamlSource,
16 download,
17 ensure_is_valid_bioimageio_yaml_name,
18)
19from ._internal.io_basics import BIOIMAGEIO_YAML, AbsoluteFilePath, FileName, ZipPath
20from ._internal.io_utils import open_bioimageio_yaml, write_yaml, write_zip
21from ._internal.packaging_context import PackagingContext
22from ._internal.url import HttpUrl
23from ._internal.utils import get_os_friendly_file_name
24from ._internal.validation_context import get_validation_context
25from ._internal.warning_levels import ERROR
26from ._io import load_description
27from .model.v0_4 import WeightsFormat
30def get_resource_package_content(
31 rd: ResourceDescr,
32 /,
33 *,
34 bioimageio_yaml_file_name: FileName = BIOIMAGEIO_YAML,
35 weights_priority_order: Optional[Sequence[WeightsFormat]] = None, # model only
36) -> Dict[FileName, Union[HttpUrl, AbsoluteFilePath, BioimageioYamlContent, ZipPath]]:
37 """
38 Args:
39 rd: resource description
40 bioimageio_yaml_file_name: RDF file name
41 # for model resources only:
42 weights_priority_order: If given, only the first weights format present in the model is included.
43 If none of the prioritized weights formats is found a ValueError is raised.
44 """
45 os_friendly_name = get_os_friendly_file_name(rd.name)
46 bioimageio_yaml_file_name = bioimageio_yaml_file_name.format(
47 name=os_friendly_name, type=rd.type
48 )
50 bioimageio_yaml_file_name = ensure_is_valid_bioimageio_yaml_name(
51 bioimageio_yaml_file_name
52 )
53 content: Dict[FileName, Union[HttpUrl, AbsoluteFilePath, ZipPath]] = {}
54 with PackagingContext(
55 bioimageio_yaml_file_name=bioimageio_yaml_file_name,
56 file_sources=content,
57 weights_priority_order=weights_priority_order,
58 ):
59 rdf_content: BioimageioYamlContent = rd.model_dump(
60 mode="json", exclude_unset=True
61 )
63 _ = rdf_content.pop("rdf_source", None)
65 return {**content, bioimageio_yaml_file_name: rdf_content}
68def _prepare_resource_package(
69 source: Union[BioimageioYamlSource, ResourceDescr],
70 /,
71 *,
72 weights_priority_order: Optional[Sequence[WeightsFormat]] = None,
73) -> Dict[FileName, Union[FilePath, BioimageioYamlContent, ZipPath]]:
74 """Prepare to package a resource description; downloads all required files.
76 Args:
77 source: A bioimage.io resource description (as file, raw YAML content or description class)
78 context: validation context
79 weights_priority_order: If given only the first weights format present in the model is included.
80 If none of the prioritized weights formats is found all are included.
81 """
82 context = get_validation_context()
83 bioimageio_yaml_file_name = context.file_name
84 if isinstance(source, ResourceDescrBase):
85 descr = source
86 elif isinstance(source, dict):
87 descr = build_description(source)
88 else:
89 opened = open_bioimageio_yaml(source)
90 bioimageio_yaml_file_name = opened.original_file_name
91 context = context.replace(
92 root=opened.original_root, file_name=opened.original_file_name
93 )
94 with context:
95 descr = build_description(opened.content)
97 if isinstance(descr, InvalidDescr):
98 raise ValueError(f"{source} is invalid: {descr.validation_summary}")
100 with context:
101 package_content = get_resource_package_content(
102 descr,
103 bioimageio_yaml_file_name=bioimageio_yaml_file_name or BIOIMAGEIO_YAML,
104 weights_priority_order=weights_priority_order,
105 )
107 local_package_content: Dict[
108 FileName, Union[FilePath, BioimageioYamlContent, ZipPath]
109 ] = {}
110 for k, v in package_content.items():
111 if not isinstance(v, (collections.abc.Mapping, ZipPath)):
112 v = download(v).path
114 local_package_content[k] = v
116 return local_package_content
119def save_bioimageio_package_as_folder(
120 source: Union[BioimageioYamlSource, ResourceDescr],
121 /,
122 *,
123 output_path: Union[NewPath, DirectoryPath, None] = None,
124 weights_priority_order: Optional[ # model only
125 Sequence[
126 Literal[
127 "keras_hdf5",
128 "onnx",
129 "pytorch_state_dict",
130 "tensorflow_js",
131 "tensorflow_saved_model_bundle",
132 "torchscript",
133 ]
134 ]
135 ] = None,
136) -> DirectoryPath:
137 """Write the content of a bioimage.io resource package to a folder.
139 Args:
140 source: bioimageio resource description
141 output_path: file path to write package to
142 weights_priority_order: If given only the first weights format present in the model is included.
143 If none of the prioritized weights formats is found all are included.
145 Returns:
146 directory path to bioimageio package folder
147 """
148 package_content = _prepare_resource_package(
149 source,
150 weights_priority_order=weights_priority_order,
151 )
152 if output_path is None:
153 output_path = Path(mkdtemp())
154 else:
155 output_path = Path(output_path)
157 output_path.mkdir(exist_ok=True, parents=True)
158 for name, src in package_content.items():
159 if isinstance(src, collections.abc.Mapping):
160 write_yaml(src, output_path / name)
161 elif isinstance(src, ZipPath):
162 extracted = Path(src.root.extract(src.name, output_path))
163 if extracted.name != src.name:
164 try:
165 shutil.move(str(extracted), output_path / src.name)
166 except Exception as e:
167 raise RuntimeError(
168 f"Failed to rename extracted file '{extracted.name}'"
169 + f" to '{src.name}'."
170 + f" (extracted from '{src.name}' in '{src.root.filename}')"
171 ) from e
172 else:
173 try:
174 shutil.copy(src, output_path / name)
175 except shutil.SameFileError:
176 pass
178 return output_path
181def save_bioimageio_package(
182 source: Union[BioimageioYamlSource, ResourceDescr],
183 /,
184 *,
185 compression: int = ZIP_DEFLATED,
186 compression_level: int = 1,
187 output_path: Union[NewPath, FilePath, None] = None,
188 weights_priority_order: Optional[ # model only
189 Sequence[
190 Literal[
191 "keras_hdf5",
192 "onnx",
193 "pytorch_state_dict",
194 "tensorflow_js",
195 "tensorflow_saved_model_bundle",
196 "torchscript",
197 ]
198 ]
199 ] = None,
200) -> FilePath:
201 """Package a bioimageio resource as a zip file.
203 Args:
204 rd: bioimageio resource description
205 compression: The numeric constant of compression method.
206 compression_level: Compression level to use when writing files to the archive.
207 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
208 output_path: file path to write package to
209 weights_priority_order: If given only the first weights format present in the model is included.
210 If none of the prioritized weights formats is found all are included.
212 Returns:
213 path to zipped bioimageio package
214 """
215 package_content = _prepare_resource_package(
216 source,
217 weights_priority_order=weights_priority_order,
218 )
219 if output_path is None:
220 output_path = Path(
221 NamedTemporaryFile(suffix=".bioimageio.zip", delete=False).name
222 )
223 else:
224 output_path = Path(output_path)
226 write_zip(
227 output_path,
228 package_content,
229 compression=compression,
230 compression_level=compression_level,
231 )
232 with get_validation_context().replace(warning_level=ERROR):
233 if isinstance((exported := load_description(output_path)), InvalidDescr):
234 raise ValueError(
235 f"Exported package '{output_path}' is invalid:"
236 + f" {exported.validation_summary}"
237 )
239 return output_path
242def save_bioimageio_package_to_stream(
243 source: Union[BioimageioYamlSource, ResourceDescr],
244 /,
245 *,
246 compression: int = ZIP_DEFLATED,
247 compression_level: int = 1,
248 output_stream: Union[IO[bytes], None] = None,
249 weights_priority_order: Optional[ # model only
250 Sequence[
251 Literal[
252 "keras_hdf5",
253 "onnx",
254 "pytorch_state_dict",
255 "tensorflow_js",
256 "tensorflow_saved_model_bundle",
257 "torchscript",
258 ]
259 ]
260 ] = None,
261) -> IO[bytes]:
262 """Package a bioimageio resource into a stream.
264 Args:
265 rd: bioimageio resource description
266 compression: The numeric constant of compression method.
267 compression_level: Compression level to use when writing files to the archive.
268 See https://docs.python.org/3/library/zipfile.html#zipfile.ZipFile
269 output_stream: stream to write package to
270 weights_priority_order: If given only the first weights format present in the model is included.
271 If none of the prioritized weights formats is found all are included.
273 Note: this function bypasses safety checks and does not load/validate the model after writing.
275 Returns:
276 stream of zipped bioimageio package
277 """
278 if output_stream is None:
279 output_stream = BytesIO()
281 package_content = _prepare_resource_package(
282 source,
283 weights_priority_order=weights_priority_order,
284 )
286 write_zip(
287 output_stream,
288 package_content,
289 compression=compression,
290 compression_level=compression_level,
291 )
293 return output_stream