Add training_data to build_model

Author: constantinpape

bioimageio/core/build_spec/build_model.py

@@ -636,6 +636,7 @@ def build_model(
     config: Optional[Dict[str, Any]] = None,
     dependencies: Optional[Union[Path, str]] = None,
     links: Optional[List[str]] = None,
+    training_data: Optional[Dict[str, str]] = None,
     root: Optional[Union[Path, str]] = None,
     add_deepimagej_config: bool = False,
     tensorflow_version: Optional[str] = None,
@@ -711,6 +712,7 @@ def build_model(
         parent: id of the parent model from which this model is derived and sha256 of the corresponding weight file.
         config: custom configuration for this model.
         dependencies: relative path to file with dependencies for this model.
+        training_data: the training data for this model, either id for a bioimageio dataset or a dataset spec.
         root: optional root path for relative paths. This can be helpful when building a spec from another model spec.
         add_deepimagej_config: add the deepimagej config to the model.
         tensorflow_version: the tensorflow version for this model. Only for tensorflow or keras weights.
@@ -887,10 +889,14 @@ def build_model(
 
     if maintainers is not None:
         kwargs["maintainers"] = [model_spec.raw_nodes.Maintainer(**m) for m in maintainers]
+
     if parent is not None:
         assert len(parent) == 2
         kwargs["parent"] = {"uri": parent[0], "sha256": parent[1]}
 
+    if training_data is not None:
+        kwargs["training_data"] = training_data
+
     try:
         model = model_spec.raw_nodes.Model(
             authors=authors,

tests/build_spec/test_build_spec.py

@@ -17,6 +17,7 @@ def _test_build_spec(
     add_deepimagej_config=False,
     use_original_covers=False,
     use_absoloute_arch_path=False,
+    training_data=None,
 ):
     from bioimageio.core.build_spec import build_model
 
@@ -115,6 +116,8 @@ def _test_build_spec(
         kwargs["pixel_sizes"] = [{"x": 5.0, "y": 5.0}]
     if use_original_covers:
         kwargs["covers"] = resolve_source(model_spec.covers, root)
+    if training_data is not None:
+        kwargs["training_data"] = training_data
 
     build_model(**kwargs)
     assert out_path.exists()
@@ -193,6 +196,21 @@ def test_build_spec_deepimagej(unet2d_nuclei_broad_model, tmp_path):
     _test_build_spec(unet2d_nuclei_broad_model, tmp_path / "model.zip", "torchscript", add_deepimagej_config=True)
 
 
+def test_build_spec_training_data1(unet2d_nuclei_broad_model, tmp_path):
+    training_data = {"id": "ilastik/stradist_dsb_training_data"}
+    _test_build_spec(unet2d_nuclei_broad_model, tmp_path / "model.zip", "torchscript", training_data=training_data)
+
+
+def test_build_spec_training_data2(unet2d_nuclei_broad_model, tmp_path):
+    training_data = {
+        "type": "dataset",
+        "name": "nucleus-training-data",
+        "description": "stardist nucleus training data",
+        "source": "https://github.com/stardist/stardist/releases/download/0.1.0/dsb2018.zip",
+    }
+    _test_build_spec(unet2d_nuclei_broad_model, tmp_path / "model.zip", "torchscript", training_data=training_data)
+
+
 def test_build_spec_deepimagej_keras(unet2d_keras, tmp_path):
     _test_build_spec(
         unet2d_keras, tmp_path / "model.zip", "keras_hdf5", add_deepimagej_config=True, tensorflow_version="1.12"