Merge pull request #65 from FynnBe/pybio

add bioimage-io module spec and package

.gitignore

@@ -2,6 +2,8 @@
 .pytest_cache
 .ipynb_checkpoints
 
+/bioimage-io/**/pytorch3dunet
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

.travis.yml

@@ -27,6 +27,7 @@ before_script:
   - conda activate test-environment
   # install test dependencies
   - conda install pytest
+  - pip install git+https://github.com/bioimage-io/python-bioimage-io#egg=pybio
 
 script:
   - pytest -v

bioimage-io/UNet3DArabidopsisOvules.model/UNet3DArabidopsisOvules.model.yaml

@@ -0,0 +1,63 @@
+name: 3D UNet Arabidopsis Ovules
+description: A 3d U-Net trained to predict the cell boundaries in confocal stacks of Arabidopsis ovules.
+cite:
+  - text: "Wolny, Adrian et al. Accurate and Versatile 3D Segmentation of Plant Tissues at Cellular Resolution. BioRxiv 2020."
+    doi: https://doi.org/10.1101/2020.01.17.910562
+authors:
+  - Adrian Wolny;@bioimage-io
+documentation: README.md
+tags: [unet3d, pytorch, arabidopsis, ovuls, cell membrane, segmentation, plant tissue]
+license: MIT
+
+format_version: 0.1.0
+language: python
+framework: pytorch
+
+source: pytorch3dunet.unet3d.model.UNet3D
+optional_kwargs:
+  in_channels: 1
+  out_channels: 1
+  layer_order: gcr  # determines the order of operators in a single layer (crg - Conv3d+ReLU+GroupNorm)
+  f_maps: [32, 64, 128, 256]   # initial number of feature maps
+  num_groups: 8  # number of groups in the groupnorm
+  final_sigmoid: true  # apply element-wise nn.Sigmoid after the final 1x1x1 convolution, otherwise apply nn.Softmax
+  is_segmentation: true  # don't touch, use postprocessing instead
+  testing: false  # don't touch, use postprocessing instead
+
+test_input: test_input.npy
+test_output: test_output.npy
+covers: [ilastik(4), ilastik(5), ilastik(6), ilastik(7), ilastik(8)]
+
+inputs:
+  - name: raw
+    axes: bczyx
+    data_type: float32
+    data_range: [-inf, inf]
+    shape: [1, 1, 144, 234, 234]
+
+outputs:
+  - name: cell_boundaries
+    axes: bczyx
+    data_type: float32
+    data_range: [0, 1]
+    halo: [0, 0, 32, 32, 32]
+    shape:
+        reference_input: raw
+        scale: [1, 1, 1, 1, 1]
+        offset: [0, 0, 0, 0, 0]
+
+prediction:
+  preprocess:
+    - spec: https://github.com/bioimage-io/pytorch-bioimage-io/blob/f71b8ac598267de88cd39e5495abd93dcda1d0a4/specs/transformations/EnsureTorch.transformation.yaml
+    - spec: https://github.com/bioimage-io/pytorch-bioimage-io/blob/f71b8ac598267de88cd39e5495abd93dcda1d0a4/specs/transformations/Cast.transformation.yaml
+      kwargs: {dtype: float32}
+    - spec: https://github.com/bioimage-io/pytorch-bioimage-io/blob/f71b8ac598267de88cd39e5495abd93dcda1d0a4/specs/transformations/NormalizeZeroMeanUnitVariance.transformation.yaml
+      kwargs: {apply_to: [0]}
+  weights:
+    source: best_model_state_dict.pth
+    hash: {md5: e12a0df1065da265aeeba81802b28f77}
+  postprocess:
+    - spec: https://github.com/bioimage-io/pytorch-bioimage-io/blob/f71b8ac598267de88cd39e5495abd93dcda1d0a4/specs/transformations/Sigmoid.transformation.yaml
+    - spec: https://github.com/bioimage-io/pytorch-bioimage-io/blob/f71b8ac598267de88cd39e5495abd93dcda1d0a4/specs/transformations/EnsureNumpy.transformation.yaml
+
+  dependencies: conda:environment.yaml

environment.yaml

@@ -15,4 +15,6 @@ dependencies:
   - scikit-learn
   - pyyaml
   - hdbscan
-  - pytest
+  - pytest
+  - pip:
+    - git+https://github.com/bioimage-io/python-bioimage-io#egg=pybio

manifest.yaml

@@ -0,0 +1,2 @@
+models:
+  - bioimage-io/UNet3DArabidopsisOvules.model/UNet3DArabidopsisOvules.model.yaml

tests/test_bioimage-io/conftest.py

@@ -0,0 +1,9 @@
+import os
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture
+def cache_path(tmp_path):
+    return Path(os.getenv("PYBIO_CACHE_PATH", tmp_path))

tests/test_bioimage-io/test_UNet3DArabidopsisOvules.py

@@ -0,0 +1,77 @@
+from io import BytesIO
+from pathlib import Path
+
+import h5py
+import numpy
+import pytest
+import torch
+
+from pybio.core.transformations import apply_transformations
+from pybio.spec import load_model
+from pybio.spec.utils import get_instance
+from pytorch3dunet.unet3d.model import UNet3D
+
+
+@pytest.fixture
+def dummy_input():
+    return [numpy.random.uniform(-2, 2, [1, 1, 144, 234, 234]).astype(numpy.float32)]
+
+
+def test_dummy_input(cache_path, dummy_input):
+    spec_path = (
+        Path(__file__).parent / "../../bioimage-io/UNet3DArabidopsisOvules.model/UNet3DArabidopsisOvules.model.yaml"
+    )
+    assert spec_path.exists()
+
+    pybio_model = load_model(str(spec_path), cache_path=cache_path)
+    for dummy, spec in zip(dummy_input, pybio_model.spec.inputs):
+        assert str(dummy.dtype) == spec.data_type
+        assert dummy.shape == spec.shape
+
+
+def test_Net3DArabidopsisOvules_forward(cache_path):
+    spec_path = (
+        Path(__file__).parent / "../../bioimage-io/UNet3DArabidopsisOvules.model/UNet3DArabidopsisOvules.model.yaml"
+    ).resolve()
+    assert spec_path.exists(), spec_path
+    pybio_model = load_model(str(spec_path), cache_path=cache_path)
+    assert pybio_model.spec.outputs[0].shape.reference_input == "raw"
+    assert pybio_model.spec.outputs[0].shape.scale == (1, 1, 1, 1, 1)
+    assert pybio_model.spec.outputs[0].shape.offset == (0, 0, 0, 0, 0)
+
+    assert isinstance(pybio_model.spec.prediction.weights.source, BytesIO)
+    assert pybio_model.spec.test_input is not None
+    assert pybio_model.spec.test_input.suffix == ".npy", pybio_model.spec.test_input.suffix
+    assert pybio_model.spec.test_output is not None
+    assert pybio_model.spec.test_output.suffix == ".npy", pybio_model.spec.test_output.suffix
+
+
+    model: torch.nn.Module = get_instance(pybio_model)
+    assert isinstance(model, UNet3D)
+    assert hasattr(model, "forward")
+    model_weights = torch.load(pybio_model.spec.prediction.weights.source, map_location=torch.device("cpu"))
+    model.load_state_dict(model_weights)
+    pre_transformations = [get_instance(trf) for trf in pybio_model.spec.prediction.preprocess]
+    post_transformations = [get_instance(trf) for trf in pybio_model.spec.prediction.postprocess]
+
+    test_ipt = numpy.load(str(pybio_model.spec.test_input))
+    assert test_ipt.shape == pybio_model.spec.inputs[0].shape
+    test_out = numpy.load(str(pybio_model.spec.test_output))
+    assert pybio_model.spec.outputs[0].shape.reference_input == pybio_model.spec.inputs[0].name
+    assert all([s == 1 for s in pybio_model.spec.outputs[0].shape.scale])
+    assert all([off == 0 for off in pybio_model.spec.outputs[0].shape.offset])
+    assert test_out.shape == pybio_model.spec.inputs[0].shape
+
+    test_roi = (slice(None), slice(None), slice(0, 32), slice(0, 32), slice(0, 32))  # to lower test mem consumption
+    ipt = apply_transformations(pre_transformations, test_ipt[test_roi])
+    assert isinstance(ipt, list)
+    assert len(ipt) == 1
+    ipt = ipt[0]
+    out = model.forward(ipt)
+    out = apply_transformations(post_transformations, out)
+    assert isinstance(out, list)
+    assert len(out) == 1
+    out = out[0]
+    # assert out.shape == pybio_model.spec.inputs[0].shape  # test_roi makes test invalid
+    assert str(out.dtype).split(".")[-1] == pybio_model.spec.outputs[0].data_type
+    assert numpy.allclose(test_out[test_roi], out, atol=0.1)  # test_roi requires atol >0.07876602

tests/test_bioimage-io/test_manifest.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+
+import pytest
+import yaml
+
+from pybio.spec import load_spec_and_kwargs, utils
+
+MANIFEST_PATH = Path(__file__).parent / "../../manifest.yaml"
+
+
+def pytest_generate_tests(metafunc):
+    if "category" in metafunc.fixturenames and "spec_path" in metafunc.fixturenames:
+        with MANIFEST_PATH.open() as f:
+            manifest = yaml.safe_load(f)
+
+        categories_and_spec_paths = [
+            (category, spec_path) for category, spec_paths in manifest.items() for spec_path in spec_paths
+        ]
+        metafunc.parametrize("category,spec_path", categories_and_spec_paths)
+
+
+@pytest.fixture
+def required_spec_kwargs():
+    repo_root = Path(__file__).parent.parent
+    kwargs = {}
+    # yaml.safe_load(
+    #     f"""
+    # """
+    # )
+
+    # testing the test data...
+    for spec_path in kwargs:
+        if not (MANIFEST_PATH.parent / spec_path).exists():
+            raise FileNotFoundError(MANIFEST_PATH.parent / spec_path)
+
+    return kwargs
+
+
+def test_load_specs_from_manifest(cache_path, category, spec_path, required_spec_kwargs):
+    kwargs = required_spec_kwargs.get(spec_path, {})
+
+    spec_path = MANIFEST_PATH.parent / spec_path
+    assert spec_path.exists()
+
+    loaded_spec = load_spec_and_kwargs(str(spec_path), **kwargs, cache_path=cache_path)
+    instance = utils.get_instance(loaded_spec)
+    assert instance