Merge pull request #233 from bioimage-io/add-training-data

Add training data

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 rdf 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,22 @@ 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"] = parent
 
+    if training_data is not None:
+        if "id" in training_data:
+            msg = f"If training data is specified via 'id' no other keys are allowed, got {training_data}"
+            assert len(training_data) == 1, msg
+            kwargs["training_data"] = training_data
+        else:
+            if "type" not in training_data:
+                training_data["type"] = "dataset"
+            if "format_version" not in training_data:
+                training_data["format_version"] = spec.dataset.format_version
+
     try:
         model = model_spec.raw_nodes.Model(
             authors=authors,

example/bioimageio-core-usage.ipynb

@@ -411,6 +411,9 @@
     "# https://github.com/bioimage-io/core-bioimage-io-python/blob/main/bioimageio/core/build_spec/build_model.py#L252\n",
     "cite = [{\"text\": cite_entry.text, \"url\": cite_entry.url} for cite_entry in model_resource.cite]\n",
     "\n",
+    "# the training data used for the model can also be specified by linking to a dataset available on bioimage.io\n",
+    "training_data = {\"id\": \"ilastik/stradist_dsb_training_data\"}\n",
+    "\n",
     "# the axes descriptions for the inputs / outputs\n",
     "input_axes = [\"bcyx\"]\n",
     "output_axes = [\"bcyx\"]\n",
@@ -441,7 +444,8 @@
     "    architecture=model_source,\n",
     "    model_kwargs=model_resource.weights[\"pytorch_state_dict\"].kwargs,\n",
     "    preprocessing=preprocessing,\n",
-    "    postprocessing=postprocessing\n",
+    "    postprocessing=postprocessing,\n",
+    "    training_data=training_data,\n",
     ")"
    ]
   },

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"