added tflite runtime support (#167)

* Release the tflite inference from tensorflow

* Update README

* Remove the bodypix_tflite from develop branch

* Add tflite inference to the tflite_inference branch

* added initial build_tflite workflow job

* added --use-feature=in-tree-build

* don't install tflite by default

* moved build_tflite up

* added tflite extra

* using dev-install-tflite

* make dev-install-tflite install build and dev depenencies

* run pytest for tflite

* using tflite extra when installing tflite

* added make dev-pytest-tflite

* linting: addressed markdown linting

* made tensorflow import optional

* added test_should_be_able_to_use_existing_tflite_model

* import tflite_runtime.interpreter

* extracted load_image

* load image using pillow

* adapted pad_and_resize_to using _pad_image_like_tensorflow

* implemented _resize_image_to_using_pillow

* added make dev-watch-tflite

* fallback to np expand_dims without tf

* extracted _get_mobilenet_preprocessed_image with np fallback

* reuse resize_image_to for scale_and_crop_to_input_tensor_shape

* automatically reduce dimension if needed

* added support for single channel in _resize_image_to_using_pillow

* extracted get_sigmoid and implemented np version

* reuse resize_image_to

* fixed failing test

* removed trailing space from requirements.txt

* cli: automatically select tflite model if full tf is not available

* don't fail with missing tf when adding alpha mask

* added tflite support to draw mask cli

* added support for remote tflite models; defined model tflite paths

* use model path constants for cli

* added TensorFlow Lite Runtime support section to readme

* ignore tflite models

* removed obsolete bodypix_tflite diectory

* fixed draw mask

* replaced pillow resize with numpy handling floats

* retain original dtype when padding

* use float32 for imagenet preprocessing

* debug logging of input image

* added list-tflite-models sub command

* fixed resnet tflite support

* added more tflite models

Co-authored-by: MrRiahi <mohammad.r.riahi@gmail.com>

Author: de-code

.github/workflows/ci.yml

@@ -21,6 +21,30 @@ jobs:
         env:
           TEST_PYPI_PASSWORD: ${{ secrets.test_pypi_password }}
 
+  build_tflite:
+    needs: []
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: [3.8]
+        include:
+        - python-version: 3.8
+
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          make venv-create SYSTEM_PYTHON=python
+          make dev-install-tflite
+      - name: Test with pytest
+        run: |
+          make dev-pytest-tflite
+
   build:
     needs: ["check_secrets"]
     runs-on: ${{ matrix.os }}

.gitignore

@@ -6,3 +6,4 @@ build
 *.egg-info
 
 *.pyc
+*.tflite

Makefile

@@ -45,13 +45,25 @@ venv-create:
 	$(SYSTEM_PYTHON) -m venv $(VENV)
 
 
-dev-install:
+dev-install-build-dependencies:
 	$(PIP) install -r requirements.build.txt
+
+
+dev-install: dev-install-build-dependencies
 	$(PIP) install \
 		-r requirements.dev.txt \
 		-r requirements.txt
 
 
+dev-install-tflite:  dev-install-build-dependencies
+	$(PIP) install -r requirements.dev.txt
+	$(PIP) install --use-feature=in-tree-build .[tflite,image]
+
+
+dev-run-pip:
+	$(PIP) $(ARGS)
+
+
 dev-venv: venv-create dev-install
 
 
@@ -75,10 +87,20 @@ dev-pytest:
 	$(PYTHON) -m pytest -p no:cacheprovider $(ARGS)
 
 
+dev-pytest-tflite:
+	$(MAKE) dev-pytest \
+		ARGS='tests/cli_test.py -k test_should_be_able_to_use_existing_tflite_model'
+
+
 dev-watch:
 	$(PYTHON) -m pytest_watch -- -p no:cacheprovider -p no:warnings $(ARGS)
 
 
+dev-watch-tflite:
+	$(MAKE) dev-watch \
+		ARGS='tests/cli_test.py -k test_should_be_able_to_use_existing_tflite_model'
+
+
 dev-test: dev-lint dev-pytest
 
 
@@ -114,6 +136,11 @@ list-models:
 		list-models
 
 
+list-tflite-models:
+	$(PYTHON) -m tf_bodypix \
+		list-tflite-models
+
+
 convert-example-draw-mask:
 	$(PYTHON) -m tf_bodypix \
 		draw-mask \
@@ -240,6 +267,66 @@ webcam-v4l2-replace-background:
 		$(ARGS)
 
 
+convert-tfjs-models-to-tflite:
+	mkdir -p "./data/tflite-models"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/050/model-stride8.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/mobilenet-float-multiplier-050-stride8-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/050/model-stride16.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/mobilenet-float-multiplier-050-stride16-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/075/model-stride8.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/mobilenet-float-multiplier-075-stride8-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/075/model-stride16.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/mobilenet-float-multiplier-075-stride16-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/100/model-stride8.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/mobilenet-float-multiplier-100-stride8-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/100/model-stride16.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/mobilenet-float-multiplier-100-stride16-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float/model-stride16.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/resnet50-float-stride16-float16.tflite"
+	$(PYTHON) -m tf_bodypix \
+			convert-to-tflite \
+			--model-path \
+			"https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float/model-stride32.json" \
+			--optimize \
+			--quantization-type=float16 \
+			--output-model-file "./data/tflite-models/resnet50-float-stride32-float16.tflite"
+
+
 docker-build:
 	docker build . -t $(IMAGE_NAME):$(IMAGE_TAG)
 

README.md

@@ -31,10 +31,11 @@ when using this project as a library:
 | ---------- | -----------
 | tf         | [TensorFlow](https://pypi.org/project/tensorflow/) (required). But you may use your own build.
 | tfjs       | TensorFlow JS Model support, using [tfjs-graph-converter](https://pypi.org/project/tfjs-graph-converter/)
+| tflite     | [tflite-runtime](https://pypi.org/project/tflite-runtime/)
 | image      | Image loading via [Pillow](https://pypi.org/project/Pillow/), required by the CLI.
 | video      | Video support via [OpenCV](https://pypi.org/project/opencv-python/)
 | webcam     | Webcam support via [OpenCV](https://pypi.org/project/opencv-python/) and [pyfakewebcam](https://pypi.org/project/pyfakewebcam/)
-| all        | All of the libraries
+| all        | All of the libraries (except `tflite-runtime`)
 
 ## Python API
 
@@ -117,6 +118,12 @@ Those URLs can be passed as the `--model-path` arguments below, or to the `downl
 
 The CLI will download and cache the model from the provided path. If no `--model-path` is provided, it will use a default model (mobilenet).
 
+To list TensorFlow Lite models instead:
+
+```bash
+python -m tf_bodypix list-tflite-models
+```
+
 ### Inputs and Outputs
 
 Most commands will work with inputs (source) and outputs.
@@ -317,7 +324,7 @@ python -m tf_bodypix \
 
 Background: [Brown Landscape Under Grey Sky](https://www.pexels.com/photo/brown-landscape-under-grey-sky-3244513/)
 
-## TensorFlow Lite support (experimental)
+## TensorFlow Lite Model support (experimental)
 
 The model path may also point to a TensorFlow Lite model (`.tflite` extension). Whether that actually improves performance may depend on the platform and available hardware.
 
@@ -330,7 +337,7 @@ python -m tf_bodypix \
     "https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/075/model-stride16.json" \
     --optimize \
     --quantization-type=float16 \
-    --output-model-file "./mobilenet-float16-stride16.tflite"
+    --output-model-file "./mobilenet-float-multiplier-075-stride16-float16.tflite"
 ```
 
 The above command is provided for convenience.
@@ -342,6 +349,12 @@ Relevant links:
 * [TF Lite post_training_quantization](https://www.tensorflow.org/lite/performance/post_training_quantization)
 * [TF GitHub #40183](https://github.com/tensorflow/tensorflow/issues/40183).
 
+## TensorFlow Lite Runtime support (experimental)
+
+This project can also be used with [tflite-runtime](https://pypi.org/project/tflite-runtime/) instead of full TensorFlow (e.g. by using the `tflite` extra).
+However, [TensorFlow Lite converter](https://www.tensorflow.org/lite/convert/) would require full TensorFlow.
+In order to avoid it, one needs to use a TensorFlow Lite model (see previous section).
+
 ## Docker Usage
 
 You could also use the Docker image if you prefer.

requirements.tflite.txt

@@ -0,0 +1 @@
+tflite-runtime==2.7.0

requirements.txt

@@ -1,4 +1,4 @@
-opencv-python==4.5.5.62 
+opencv-python==4.5.5.62
 Pillow==8.4.0; python_version < "3.7"
 Pillow==9.0.1; python_version >= "3.7"
 pyfakewebcam==0.1.0

setup.py

@@ -12,6 +12,10 @@
     REQUIRED_PACKAGES = f.readlines()
 
 
+with open('requirements.tflite.txt', 'r', encoding='utf-8') as f:
+    TFLITE_REQUIRED_PACKAGES = f.readlines()
+
+
 with open('README.md', 'r', encoding='utf-8') as f:
     LONG_DESCRIPTION = '\n'.join([
         line.rstrip()
@@ -30,6 +34,10 @@ def local_scheme(version):
     get_requirements_with_groups(REQUIRED_PACKAGES)
 )
 
+ALL_EXTRAS = {
+    **EXTRAS,
+    'tflite': TFLITE_REQUIRED_PACKAGES
+}
 
 packages = find_packages(exclude=["tests", "tests.*"])
 
@@ -42,7 +50,7 @@ def local_scheme(version):
     author="Daniel Ecer",
     url="https://github.com/de-code/python-tf-bodypix",
     install_requires=DEFAULT_REQUIRED_PACKAGES,
-    extras_require=EXTRAS,
+    extras_require=ALL_EXTRAS,
     packages=packages,
     include_package_data=True,
     description='Python implemention of the TensorFlow BodyPix model.',

tests/cli_test.py

@@ -1,9 +1,9 @@
 import logging
 from pathlib import Path
 
-from tf_bodypix.download import BodyPixModelPaths
+from tf_bodypix.download import ALL_TENSORFLOW_LITE_BODYPIX_MODEL_PATHS, BodyPixModelPaths
 from tf_bodypix.model import ModelArchitectureNames
-from tf_bodypix.cli import main
+from tf_bodypix.cli import DEFAULT_MODEL_TFLITE_PATH, main
 
 
 LOGGER = logging.getLogger(__name__)
@@ -97,6 +97,15 @@ def test_should_list_all_default_model_urls(self, capsys):
         missing_urls = set(expected_urls) - set(output_urls)
         assert not missing_urls
 
+    def test_should_list_all_default_tflite_models(self, capsys):
+        expected_urls = ALL_TENSORFLOW_LITE_BODYPIX_MODEL_PATHS
+        main(['list-tflite-models'])
+        captured = capsys.readouterr()
+        output_urls = captured.out.splitlines()
+        LOGGER.debug('output_urls: %s', output_urls)
+        missing_urls = set(expected_urls) - set(output_urls)
+        assert not missing_urls
+
     def test_should_be_able_to_convert_to_tflite_and_use_model(self, temp_dir: Path):
         output_model_file = temp_dir / 'model.tflite'
         main([
@@ -115,3 +124,14 @@ def test_should_be_able_to_convert_to_tflite_and_use_model(self, temp_dir: Path)
             '--source=%s' % EXAMPLE_IMAGE_URL,
             '--output=%s' % output_image_path
         ])
+
+    def test_should_be_able_to_use_existing_tflite_model(self, temp_dir: Path):
+        output_image_path = temp_dir / 'mask.jpg'
+        main([
+            'draw-mask',
+            '--model-path=%s' % DEFAULT_MODEL_TFLITE_PATH,
+            '--model-architecture=%s' % ModelArchitectureNames.MOBILENET_V1,
+            '--output-stride=16',
+            '--source=%s' % EXAMPLE_IMAGE_URL,
+            '--output=%s' % output_image_path
+        ])

tf_bodypix/bodypix_js_utils/decode_part_map.py

@@ -1,16 +1,26 @@
 # based on:
 # https://github.com/tensorflow/tfjs-models/blob/body-pix-v2.0.4/body-pix/src/decode_part_map.ts
 
-import tensorflow as tf
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
 
 import numpy as np
 
 
+DEFAULT_DTYPE = (
+    tf.int32 if tf is not None else np.int32
+)
+
+
 def to_mask_tensor(
     segment_scores: np.ndarray,
     threshold: float,
-    dtype: type = tf.int32
+    dtype: type = DEFAULT_DTYPE
 ) -> np.ndarray:
+    if tf is None:
+        return (segment_scores > threshold).astype(dtype)
     return tf.cast(
         tf.greater(segment_scores, threshold),
         dtype