annotate relative and absolute tolerance

Author: FynnBe

bioimageio/core/_resource_tests.py

@@ -37,6 +37,7 @@
 from bioimageio.spec._internal.common_nodes import ResourceDescrBase
 from bioimageio.spec._internal.io import is_yaml_value
 from bioimageio.spec._internal.io_utils import read_yaml, write_yaml
+from bioimageio.spec._internal.types import AbsoluteTolerance, RelativeTolerance
 from bioimageio.spec._internal.validation_context import validation_context_var
 from bioimageio.spec.common import BioimageioYamlContent, PermissiveFileSource, Sha256
 from bioimageio.spec.model import v0_4, v0_5
@@ -120,8 +121,8 @@ def test_model(
     source: Union[v0_4.ModelDescr, v0_5.ModelDescr, PermissiveFileSource],
     weight_format: Optional[WeightsFormat] = None,
     devices: Optional[List[str]] = None,
-    absolute_tolerance: float = 1.5e-4,
-    relative_tolerance: float = 1e-4,
+    absolute_tolerance: AbsoluteTolerance = 1.5e-4,
+    relative_tolerance: RelativeTolerance = 1e-4,
     decimal: Optional[int] = None,
     *,
     determinism: Literal["seed_only", "full"] = "seed_only",
@@ -152,8 +153,8 @@ def test_description(
     format_version: Union[Literal["discover", "latest"], str] = "discover",
     weight_format: Optional[WeightsFormat] = None,
     devices: Optional[Sequence[str]] = None,
-    absolute_tolerance: float = 1.5e-4,
-    relative_tolerance: float = 1e-4,
+    absolute_tolerance: AbsoluteTolerance = 1.5e-4,
+    relative_tolerance: RelativeTolerance = 1e-4,
     decimal: Optional[int] = None,
     determinism: Literal["seed_only", "full"] = "seed_only",
     expected_type: Optional[str] = None,
@@ -236,8 +237,8 @@ def _test_in_env(
     weight_format: Optional[WeightsFormat],
     conda_env: Optional[BioimageioCondaEnv],
     devices: Optional[Sequence[str]],
-    absolute_tolerance: float,
-    relative_tolerance: float,
+    absolute_tolerance: AbsoluteTolerance,
+    relative_tolerance: RelativeTolerance,
     determinism: Literal["seed_only", "full"],
     run_command: Callable[[Sequence[str]], None],
 ) -> ValidationSummary:
@@ -354,8 +355,8 @@ def load_description_and_test(
     format_version: Union[Literal["discover", "latest"], str] = "discover",
     weight_format: Optional[WeightsFormat] = None,
     devices: Optional[Sequence[str]] = None,
-    absolute_tolerance: float = 1.5e-4,
-    relative_tolerance: float = 1e-4,
+    absolute_tolerance: AbsoluteTolerance = 1.5e-4,
+    relative_tolerance: RelativeTolerance = 1e-4,
     decimal: Optional[int] = None,
     determinism: Literal["seed_only", "full"] = "seed_only",
     expected_type: Optional[str] = None,

bioimageio/core/weight_converters/_add_weights.py

@@ -1,23 +1,35 @@
-from abc import ABC
-from typing import Optional, Sequence, Union, assert_never, final
+from copy import deepcopy
+from pathlib import Path
+from typing import List, Optional, Sequence, Union
 
 from bioimageio.spec.model import v0_4, v0_5
 
 
 def increase_available_weight_formats(
-    model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr],
-    source_format: v0_5.WeightsFormat,
-    target_format: v0_5.WeightsFormat,
+    model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr],
     *,
+    source_format: Optional[v0_5.WeightsFormat] = None,
+    target_format: Optional[v0_5.WeightsFormat] = None,
+    output_path: Path,
     devices: Optional[Sequence[str]] = None,
-):
-    if not isinstance(model_description, (v0_4.ModelDescr, v0_5.ModelDescr)):
+) -> Union[v0_4.ModelDescr, v0_5.ModelDescr]:
+    """Convert neural network weights to other formats and add them to the model description"""
+    if not isinstance(model_descr, (v0_4.ModelDescr, v0_5.ModelDescr)):
         raise TypeError(
-            f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_description)}"
+            f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_descr)}"
         )
 
-    if (source_format, target_format) == ("pytorch_state_dict", "onnx"):
-        from .pytorch_to_onnx import convert_pytorch_to_onnx
+    if source_format is None:
+        available = [wf for wf, w in model_descr.weights if w is not None]
+        missing = [wf for wf, w in model_descr.weights if w is None]
+    else:
+        available = [source_format]
+        missing = [target_format]
+
+    if "pytorch_state_dict" in available and "onnx" in missing:
+        from .pytorch_to_onnx import convert
+
+        onnx = convert(model_descr)
 
     else:
         raise NotImplementedError(