Fix more issues in build_spec and update example notebook

Author: constantinpape

bioimageio/core/build_spec/build_model.py

@@ -59,23 +59,24 @@ def _get_pytorch_state_dict_weight_kwargs(architecture, model_kwargs, root):
     assert architecture is not None
     tmp_archtecture = None
     weight_kwargs = {"kwargs": model_kwargs} if model_kwargs else {}
-    arch = spec.shared.fields.ImportableSource().deserialize(architecture)
-    if isinstance(arch, ImportableSourceFile):
-        if os.path.isabs(arch.source_file):
-            tmp_archtecture = Path("this_model_architecture.py")
+    if ":" in architecture:
+        arch_file, callable_name = architecture.replace("::", ":").split(":")
 
-            copyfile(arch.source_file, root / tmp_archtecture)
-            arch = ImportableSourceFile(arch.callable_name, tmp_archtecture)
+        # this goes haywire if we pass an absolute path, so need to copt to a tmp relative path
+        if os.path.isabs(arch_file):
+            tmp_archtecture = Path("this_model_architecture.py")
+            copyfile(arch_file, root / tmp_archtecture)
+            arch = ImportableSourceFile(callable_name, tmp_archtecture)
+        else:
+            arch = ImportableSourceFile(callable_name, Path(arch_file))
 
         arch_hash = _get_hash(root / arch.source_file)
         weight_kwargs["architecture_sha256"] = arch_hash
-    elif isinstance(arch, ImportableModule):
-        pass
     else:
-        raise NotImplementedError(arch)
+        arch = spec.shared.fields.ImportableSource().deserialize(architecture)
+        assert isinstance(arch, ImportableModule)
 
     weight_kwargs["architecture"] = arch
-
     return weight_kwargs, tmp_archtecture
 
 

example/.gitignore

@@ -0,0 +1 @@
+.ipynb_checkpoints/

example/bioimageio-core-usage.ipynb

@@ -144,6 +144,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# the function 'test_model' from 'bioimageio.core.resource_tests' can be used to fully test the model,\n",
+    "# including running prediction for the test input(s) and checking that they agree with the test output(s)\n",
+    "# before using a model, it is recommended to check that it properly works with this function\n",
+    "# 'test_model' returns a dict, if there are any errros they will be in the key \"error\"\n",
+    "# if the model passes it will be None\n",
     "from bioimageio.core.resource_tests import test_model\n",
     "test_result = test_model(model_resource)\n",
     "if test_result[\"error\"]:\n",
@@ -326,7 +331,7 @@
     "\n",
     "`bioimageio.core` also implements functionality to create a model package compatible with the [bioimageio model spec](https://github.com/bioimage-io/spec-bioimage-io/blob/gh-pages/model_spec_latest.md) ready to be shared via\n",
     "the [bioimage.io model zoo](https://bioimage.io/#/).\n",
-    "Here, we will use this functionality to create two models, one that adds thresholding of the outputs as post-processing and another that also adds weights in torchscript format."
+    "Here, we will use this functionality to create two models, one that adds thresholding as post-processing to the outputs and another one that also adds weights in torchscript format."
    ]
   },
   {
@@ -336,20 +341,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# get the python file defining the architecture\n",
+    "# get the python file defining the architecture.\n",
+    "# this is only required for models with pytorch_state_dict weights\n",
     "def get_architecture_source(rdf):\n",
     "    # here, we need the raw resource, which contains the information from the resource description\n",
     "    # before evaluation, e.g. the file and name of the python file with the model architecture\n",
     "    raw_resource = bioimageio.core.load_raw_resource_description(rdf)\n",
-    "    # the source information\n",
-    "    model_source = raw_resource.source\n",
+    "    # the python file defining the architecture for the pytorch weihgts\n",
+    "    model_source = raw_resource.weights[\"pytorch_state_dict\"].architecture\n",
     "    # download the source file if necessary\n",
     "    source_file = bioimageio.core.resource_io.utils.resolve_source(\n",
     "        model_source.source_file\n",
     "    )\n",
     "    # if the source file path does not exist, try combining it with the root path of the model\n",
     "    if not os.path.exists(source_file):\n",
-    "        source_file = os.path.join(raw_resource.root_path, source_file)\n",
+    "        source_file = os.path.join(raw_resource.root_path, os.path.split(source_file)[1])\n",
     "    assert os.path.exists(source_file), source_file\n",
     "    class_name = model_source.callable_name\n",
     "    return f\"{source_file}:{class_name}\""
@@ -390,9 +396,9 @@
     "# the nucleus segmentation model we have obtained from bioimage.io\n",
     "# this field is optional and only needs to be given for models that are created based on\n",
     "# other models from bioimage.io\n",
-    "# the parent is specified via it's doi and the hash of the weight file\n",
-    "weight_file = model_resource.weights[\"pytorch_state_dict\"].source\n",
-    "with open(weight_file, \"rb\") as f:\n",
+    "# the parent is specified via it's doi and the hash of its rdf file\n",
+    "rdf_file = os.path.join(model_resource.root_path, \"rdf.yaml\")\n",
+    "with open(rdf_file, \"rb\") as f:\n",
     "    weight_hash = hashlib.sha256(f.read()).hexdigest()\n",
     "parent = (rdf_doi, weight_hash)\n",
     "\n",
@@ -405,10 +411,24 @@
     "# for more informantion, check out the function signature\n",
     "# https://github.com/bioimage-io/core-bioimage-io-python/blob/main/bioimageio/core/build_spec/build_model.py#L252\n",
     "cite = {cite_entry.text: cite_entry.url for cite_entry in model_resource.cite}\n",
+    "\n",
+    "# the axes descriptions for the inputs / outputs\n",
+    "input_axes = [\"bcyx\"]\n",
+    "output_axes = [\"bcyx\"]\n",
+    "\n",
+    "# the pytorch_state_dict weight file\n",
+    "weight_file = model_resource.weights[\"pytorch_state_dict\"].source\n",
+    "\n",
+    "# the path to save the new model with torchscript weights\n",
+    "zip_path = f\"{model_root}/new_model2.zip\"\n",
+    "\n",
+    "# build the model! it will be saved to 'zip_path'\n",
     "new_model_raw = build_model(\n",
-    "    weight_file,\n",
+    "    weight_uri=weight_file,\n",
     "    test_inputs=model_resource.test_inputs,\n",
     "    test_outputs=[new_output_path],\n",
+    "    input_axes=input_axes,\n",
+    "    output_axes=output_axes,\n",
     "    output_path=zip_path,\n",
     "    name=name,\n",
     "    description=\"nucleus segmentation model with thresholding\",\n",
@@ -419,9 +439,8 @@
     "    tags=[\"nucleus-segmentation\"],\n",
     "    cite=cite,\n",
     "    parent=parent,\n",
-    "    root=model_root,\n",
-    "    source=model_source,\n",
-    "    model_kwargs=model_resource.kwargs,\n",
+    "    architecture=model_source,\n",
+    "    model_kwargs=model_resource.weights[\"pytorch_state_dict\"].kwargs,\n",
     "    preprocessing=preprocessing,\n",
     "    postprocessing=postprocessing\n",
     ")"
@@ -456,17 +475,17 @@
    "outputs": [],
    "source": [
     "# `convert_weigths_to_pytorch_script` creates torchscript weigths based on the weights loaded from pytorch_state_dict\n",
-    "from bioimageio.core.weight_converter.torch import convert_weights_to_pytorch_script\n",
+    "from bioimageio.core.weight_converter.torch import convert_weights_to_torchscript\n",
     "# `add_weights` adds new weights to the model specification\n",
     "from bioimageio.core.build_spec import add_weights\n",
     "\n",
     "# the path to save the newly created torchscript weights\n",
     "weight_path = os.path.join(model_root, \"weights.torchscript\")\n",
-    "convert_weights_to_pytorch_script(new_model, weight_path)\n",
+    "convert_weights_to_torchscript(new_model, weight_path)\n",
     "\n",
     "# the path to save the new model with torchscript weights\n",
     "zip_path = f\"{model_root}/new_model2.zip\"\n",
-    "new_model2_raw = add_weights(new_model_raw, weight_path, weight_type=\"pytorch_script\", output_path=zip_path)"
+    "new_model2_raw = add_weights(new_model_raw, weight_path, weight_type=\"torchscript\", output_path=zip_path)"
    ]
   },
   {
@@ -478,7 +497,7 @@
    "source": [
     "# load the new model from the zipped package, run prediction and check the result\n",
     "new_model = bioimageio.core.load_resource_description(zip_path)\n",
-    "prediction = predict_numpy(new_model, input_image, weight_format=\"pytorch_script\")\n",
+    "prediction = predict_numpy(new_model, input_image, weight_format=\"torchscript\")\n",
     "show_images(input_image, prediction, names=[\"input\", \"binarized-prediction\"])"
    ]
   },