Merge pull request #304 from bioimage-io/tf-2

Add support for tensorflow 2

Author: constantinpape

.github/workflows/build.yml

@@ -69,11 +69,11 @@ jobs:
       shell: bash -l {0}
       run: pytest --disable-pytest-warnings
 
-  test-base-bioimage-spec-tf-legacy:
+  test-base-bioimage-spec-tf:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7]
+        python-version: [3.7, 3.8, 3.9]
     steps:
     - uses: actions/checkout@v2
     - name: install dependencies
@@ -92,6 +92,33 @@ jobs:
         conda remove --force bioimageio.spec
         pip install --no-deps git+https://github.com/bioimage-io/spec-bioimage-io
         pip install --no-deps -e .
+    - name: pytest-base-bioimage-spec-tf
+      shell: bash -l {0}
+      run: pytest --disable-pytest-warnings
+
+  test-base-bioimage-spec-tf-legacy:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.7]
+    steps:
+    - uses: actions/checkout@v2
+    - name: install dependencies
+      uses: conda-incubator/setup-miniconda@v2
+      with:
+        auto-update-conda: true
+        # we need mamba to resolve environment-tf
+        mamba-version: "*"
+        channel-priority: flexible
+        activate-environment: bio-core-tf-legacy
+        environment-file: dev/environment-tf-legacy.yaml
+        python-version: ${{ matrix.python-version }}
+    - name: additional setup
+      shell: bash -l {0}
+      run: |
+        conda remove --force bioimageio.spec
+        pip install --no-deps git+https://github.com/bioimage-io/spec-bioimage-io
+        pip install --no-deps -e .
     - name: pytest-base-bioimage-spec-tf-legacy
       shell: bash -l {0}
       run: pytest --disable-pytest-warnings

bioimageio/core/prediction_pipeline/_model_adapters/_tensorflow_model_adapter.py

@@ -97,9 +97,14 @@ def _forward_tf(self, *input_tensors):
 
         return res
 
-    def _forward_keras(self, input_tensors):
+    def _forward_keras(self, *input_tensors):
         tf_tensor = [tf.convert_to_tensor(ipt) for ipt in input_tensors]
-        result = self._model.forward(*tf_tensor)
+
+        try:
+            result = self._model.forward(*tf_tensor)
+        except AttributeError:
+            result = self._model.predict(*tf_tensor)
+
         if not isinstance(result, (tuple, list)):
             result = [result]
 

bioimageio/core/weight_converter/keras/tensorflow.py

@@ -11,6 +11,28 @@
 from tensorflow import saved_model
 
 
+def _zip_weights(output_path):
+    zipped_model = f"{output_path}.zip"
+    # zip the weights
+    file_paths = []
+    for folder_names, subfolder, filenames in os.walk(os.path.join(output_path)):
+        for filename in filenames:
+            # create complete filepath of file in directory
+            file_paths.append(os.path.join(folder_names, filename))
+
+    with ZipFile(zipped_model, "w") as zip_obj:
+        for f in file_paths:
+            # Add file to zip
+            zip_obj.write(f, os.path.relpath(f, output_path))
+
+    try:
+        shutil.rmtree(output_path)
+    except Exception:
+        print("TensorFlow bundled model was not removed after compression")
+
+    return zipped_model
+
+
 # adapted from
 # https://github.com/deepimagej/pydeepimagej/blob/master/pydeepimagej/yaml/create_config.py#L236
 def _convert_tf1(keras_weight_path, output_path, input_name, output_name, zip_weights):
@@ -40,26 +62,27 @@ def build_tf_model():
         build_tf_model()
 
     if zip_weights:
-        zipped_model = f"{output_path}.zip"
-        # zip the weights
-        file_paths = []
-        for folder_names, subfolder, filenames in os.walk(os.path.join(output_path)):
-            for filename in filenames:
-                # create complete filepath of file in directory
-                file_paths.append(os.path.join(folder_names, filename))
-
-        with ZipFile(zipped_model, "w") as zip_obj:
-            for f in file_paths:
-                # Add file to zip
-                zip_obj.write(f, os.path.relpath(f, output_path))
-
-        try:
-            shutil.rmtree(output_path)
-        except Exception:
-            print("TensorFlow bundled model was not removed after compression")
-        print("TensorFlow model exported to", zipped_model)
-    else:
-        print("TensorFlow model exported to", output_path)
+        output_path = _zip_weights(output_path)
+    print("TensorFlow model exported to", output_path)
+
+    return 0
+
+
+def _convert_tf2(keras_weight_path, output_path, zip_weights):
+    try:
+        # try to build the tf model with the keras import from tensorflow
+        from tensorflow import keras
+    except Exception:
+        # if the above fails try to export with the standalone keras
+        import keras
+
+    model = keras.models.load_model(keras_weight_path)
+    keras.models.save_model(model, output_path)
+
+    if zip_weights:
+        output_path = _zip_weights(output_path)
+    print("TensorFlow model exported to", output_path)
+
     return 0
 
 
@@ -104,4 +127,4 @@ def convert_weights_to_tensorflow_saved_model_bundle(
             )
         return _convert_tf1(weight_path, str(path_), model.inputs[0].name, model.outputs[0].name, zip_weights)
     else:
-        raise NotImplementedError("Weight conversion for tensorflow 2 is not yet implemented.")
+        return _convert_tf2(weight_path, str(path_), zip_weights)

dev/environment-tf-legacy.yaml

@@ -0,0 +1,17 @@
+name: bio-core-tf-legacy
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - black
+  - bioimageio.spec
+  - conda-build
+  - h5py >=2.10,<2.11
+  - mypy
+  - pip
+  - pytest
+  - python >=3.7,<3.8  # this environment is only available for python 3.7
+  - xarray
+  - tensorflow >1.14,<2.0
+  - tifffile <=2022.4.8  # pin fixes Syntax error; see https://github.com/bioimage-io/core-bioimage-io-python/pull/259
+  - keras

dev/environment-tf.yaml

@@ -6,12 +6,10 @@ dependencies:
   - black
   - bioimageio.spec
   - conda-build
-  - h5py >=2.10,<2.11
   - mypy
   - pip
   - pytest
-  - python >=3.7,<3.8  # this environment is only available for python 3.7
+  - python
   - xarray
-  - tensorflow >1.14,<2.0
+  - tensorflow >=2.9,<3.0
   - tifffile <=2022.4.8  # pin fixes Syntax error; see https://github.com/bioimage-io/core-bioimage-io-python/pull/259
-  - keras

tests/build_spec/test_build_spec.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from marshmallow import missing
 
 import bioimageio.spec as spec
@@ -6,6 +8,13 @@
 from bioimageio.core.resource_io.utils import resolve_source
 from bioimageio.core.resource_tests import test_model as _test_model
 
+try:
+    import tensorflow
+except ImportError:
+    tf_version = None
+else:
+    tf_version: Optional[str] = ".".join(tensorflow.__version__.split(".")[:2])
+
 
 def _test_build_spec(
     spec_path,
@@ -175,18 +184,18 @@ def test_build_spec_onnx(any_onnx_model, tmp_path):
 
 def test_build_spec_keras(any_keras_model, tmp_path):
     _test_build_spec(
-        any_keras_model, tmp_path / "model.zip", "keras_hdf5", tensorflow_version="1.12"
+        any_keras_model, tmp_path / "model.zip", "keras_hdf5", tensorflow_version=tf_version
     )  # todo: keras for tf 2??
 
 
 def test_build_spec_tf(any_tensorflow_model, tmp_path):
     _test_build_spec(
-        any_tensorflow_model, tmp_path / "model.zip", "tensorflow_saved_model_bundle", tensorflow_version="1.12"
+        any_tensorflow_model, tmp_path / "model.zip", "tensorflow_saved_model_bundle", tensorflow_version=tf_version
     )  # check tf version
 
 
 def test_build_spec_tfjs(any_tensorflow_js_model, tmp_path):
-    _test_build_spec(any_tensorflow_js_model, tmp_path / "model.zip", "tensorflow_js", tensorflow_version="1.12")
+    _test_build_spec(any_tensorflow_js_model, tmp_path / "model.zip", "tensorflow_js", tensorflow_version=tf_version)
 
 
 def test_build_spec_deepimagej(unet2d_nuclei_broad_model, tmp_path):
@@ -220,7 +229,7 @@ def test_build_spec_parent2(unet2d_nuclei_broad_model, tmp_path):
 
 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"
+        unet2d_keras, tmp_path / "model.zip", "keras_hdf5", add_deepimagej_config=True, tensorflow_version=tf_version
     )
 
 

tests/conftest.py

@@ -19,15 +19,20 @@
 torchscript_models = ["unet2d_multi_tensor", "unet2d_nuclei_broad_model"]
 onnx_models = ["unet2d_multi_tensor", "unet2d_nuclei_broad_model", "hpa_densenet"]
 tensorflow1_models = ["stardist"]
-tensorflow2_models = []
-keras_models = ["unet2d_keras"]
+tensorflow2_models = ["unet2d_keras_tf2"]
+keras_tf1_models = ["unet2d_keras"]
+keras_tf2_models = ["unet2d_keras_tf2"]
 tensorflow_js_models = []
 
 model_sources = {
     "unet2d_keras": (
         "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/"
         "unet2d_keras_tf/rdf.yaml"
     ),
+    "unet2d_keras_tf2": (
+        "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/"
+        "unet2d_keras_tf2/rdf.yaml"
+    ),
     "unet2d_nuclei_broad_model": (
         "https://raw.githubusercontent.com/bioimage-io/spec-bioimage-io/main/example_specs/models/"
         "unet2d_nuclei_broad/rdf.yaml"
@@ -91,12 +96,6 @@
 skip_tensorflow = tensorflow is None
 skip_tensorflow_js = True  # TODO: add a tensorflow_js example model
 
-try:
-    import keras
-except ImportError:
-    keras = None
-skip_keras = keras is None
-
 # load all model packages we need for testing
 load_model_packages = set()
 if not skip_torch:
@@ -106,13 +105,14 @@
     load_model_packages |= set(onnx_models)
 
 if not skip_tensorflow:
-    load_model_packages |= set(keras_models)
     load_model_packages |= set(tensorflow_js_models)
     if tf_major_version == 1:
+        load_model_packages |= set(keras_tf1_models)
         load_model_packages |= set(tensorflow1_models)
         load_model_packages.add("stardist_wrong_shape")
         load_model_packages.add("stardist_wrong_shape2")
     elif tf_major_version == 2:
+        load_model_packages |= set(keras_tf2_models)
         load_model_packages |= set(tensorflow2_models)
 
 
@@ -146,14 +146,12 @@ def any_onnx_model(request):
     return pytest.model_packages[request.param]
 
 
-@pytest.fixture(params=[] if skip_tensorflow else (set(tensorflow1_models) | set(tensorflow2_models)))
+@pytest.fixture(params=[] if skip_tensorflow else tensorflow1_models if tf_major_version == 1 else tensorflow2_models)
 def any_tensorflow_model(request):
-    name = request.param
-    if (tf_major_version == 1 and name in tensorflow1_models) or (tf_major_version == 2 and name in tensorflow2_models):
-        return pytest.model_packages[name]
+    return pytest.model_packages[request.param]
 
 
-@pytest.fixture(params=[] if skip_keras else keras_models)
+@pytest.fixture(params=[] if skip_tensorflow else keras_tf1_models if tf_major_version == 1 else keras_tf2_models)
 def any_keras_model(request):
     return pytest.model_packages[request.param]
 
@@ -178,21 +176,17 @@ def any_model(request):
 #
 
 
-@pytest.fixture(
-    params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_fixed_shape"]
-)
+@pytest.fixture(params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_fixed_shape"])
 def unet2d_fixed_shape_or_not(request):
     return pytest.model_packages[request.param]
 
 
-@pytest.fixture(
-    params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_multi_tensor"]
-)
+@pytest.fixture(params=[] if skip_torch else ["unet2d_nuclei_broad_model", "unet2d_multi_tensor"])
 def unet2d_multi_tensor_or_not(request):
     return pytest.model_packages[request.param]
 
 
-@pytest.fixture(params=[] if skip_keras else ["unet2d_keras"])
+@pytest.fixture(params=[] if skip_tensorflow else ["unet2d_keras" if tf_major_version == 1 else "unet2d_keras_tf2"])
 def unet2d_keras(request):
     return pytest.model_packages[request.param]