Add utility function for adding additional weight formats to model spec

bioimageio/spec/latest/__init__.py

@@ -1,2 +1,2 @@
 from bioimageio.spec.v0_3 import *  # noqa
-from .build_spec import build_spec
+from .build_spec import build_spec, add_weights

bioimageio/spec/latest/build_spec.py

@@ -48,7 +48,6 @@ def _infer_weight_type(path):
         raise ValueError(f"Could not infer weight type from extension {ext} for weight file {path}")
 
 
-# TODO extend supported weight types
 def _get_weights(weight_uri, weight_type, source, root, **kwargs):
     weight_path = _get_local_path(weight_uri, root)
     if weight_type is None:
@@ -219,6 +218,7 @@ def _build_cite(cite):
     return citation_list
 
 
+# TODO we should make the name more specific: "build_model_spec"?
 # NOTE does not support multiple input / output tensors yet
 # to implement this we should wait for 0.4.0, see also
 # https://github.com/bioimage-io/spec-bioimage-io/issues/70#issuecomment-825737433
@@ -412,3 +412,23 @@ def build_spec(
     model = spec.schema.Model().load(serialized)
 
     return model
+
+
+def add_weights(
+    model: spec.raw_nodes.Model,
+    weight_uri: str,
+    root: Optional[str] = None,
+    weight_type: Optional[str] = None,
+    **weight_kwargs
+):
+    """ Add weight entry to bioimage.io model.
+    """
+    new_weights = _get_weights(weight_uri, weight_type, None, root, **weight_kwargs)[0]
+    model.weights.update(new_weights)
+
+    # FIXME this fails with
+    # ImportedSource(factory=<class 'user_imports.5e008e787272408180a19fd72b83134b.UNet2d'>) has unexpected type <class 'bioimageio.spec.shared.nodes.ImportedSource'>
+    serialized = spec.schema.Model().dump(model)
+    model = spec.schema.Model().load(serialized)
+
+    return model

bioimageio/spec/v0_3/utils.py

@@ -35,9 +35,9 @@ def _(source: dict, root_path: Optional[pathlib.Path] = None) -> Tuple[raw_nodes
 def _(source: os.PathLike, root_path: Optional[pathlib.Path] = None) -> Tuple[raw_nodes.Model, pathlib.Path]:
     source = pathlib.Path(source)
 
-    suffixes = source.suffixes
-    if len(suffixes) < 2 or suffixes[-1] not in (".yml", ".yaml") or source.suffixes[-2] != ".model":
-        raise ValidationError(f"invalid suffixes {''.join(suffixes)} for source {source}")
+    # suffixes = source.suffixes
+    # if len(suffixes) < 2 or suffixes[-1] not in (".yml", ".yaml") or source.suffixes[-2] != ".model":
+    #     raise ValidationError(f"invalid suffixes {''.join(suffixes)} for source {source}")
 
     data = yaml.load(source)