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