Add pip package builder for model_maker

PiperOrigin-RevId: 493989013

mediapipe/model_maker/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include pip_src/mediapipe_model_maker/models *

mediapipe/model_maker/__init__.py

@@ -11,3 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+
+from mediapipe.model_maker.python.core.utils import quantization
+from mediapipe.model_maker.python.vision import image_classifier
+from mediapipe.model_maker.python.vision import gesture_recognizer
+from mediapipe.model_maker.python.text import text_classifier

mediapipe/model_maker/python/core/utils/file_util.py

@@ -19,7 +19,7 @@ import os
 
 
 def get_absolute_path(file_path: str) -> str:
-  """Gets the absolute path of a file.
+  """Gets the absolute path of a file in the model_maker directory.
 
   Args:
     file_path: The path to a file relative to the `mediapipe` dir
@@ -27,10 +27,17 @@ def get_absolute_path(file_path: str) -> str:
   Returns:
    The full path of the file
   """
-  # Extract the file path before mediapipe/ as the `base_dir`. By joining it
-  # with the `path` which defines the relative path under mediapipe/, it
-  # yields to the absolute path of the model files directory.
+  # Extract the file path before and including 'model_maker' as the
+  # `mm_base_dir`. By joining it with the `path` after 'model_maker/', it
+  # yields to the absolute path of the model files directory. We must join
+  # on 'model_maker' because in the pypi package, the 'model_maker' directory
+  # is renamed to 'mediapipe_model_maker'. So we have to join on model_maker
+  # to ensure that the `mm_base_dir` path includes the renamed
+  # 'mediapipe_model_maker' directory.
   cwd = os.path.dirname(__file__)
-  base_dir = cwd[:cwd.rfind('mediapipe')]
-  absolute_path = os.path.join(base_dir, file_path)
+  cwd_stop_idx = cwd.rfind('model_maker') + len('model_maker')
+  mm_base_dir = cwd[:cwd_stop_idx]
+  file_path_start_idx = file_path.find('model_maker') + len('model_maker') + 1
+  mm_relative_path = file_path[file_path_start_idx:]
+  absolute_path = os.path.join(mm_base_dir, mm_relative_path)
   return absolute_path

mediapipe/model_maker/python/text/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2022 The MediaPipe Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

mediapipe/model_maker/requirements.txt

@@ -1,6 +1,8 @@
 absl-py
+mediapipe==0.9.1
 numpy
-opencv-contrib-python
-tensorflow
+opencv-python
+tensorflow>=2.10
 tensorflow-datasets
 tensorflow-hub
+tf-models-official>=2.10.1

mediapipe/model_maker/setup.py

@@ -0,0 +1,147 @@
+"""Copyright 2020-2022 The MediaPipe Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Setup for Mediapipe-Model-Maker package with setuptools.
+"""
+
+import glob
+import os
+import shutil
+import subprocess
+import sys
+import setuptools
+
+
+__version__ = 'dev'
+MM_ROOT_PATH = os.path.dirname(os.path.abspath(__file__))
+# Build dir to copy all necessary files and build package
+SRC_NAME = 'pip_src'
+BUILD_DIR = os.path.join(MM_ROOT_PATH, SRC_NAME)
+BUILD_MM_DIR = os.path.join(BUILD_DIR, 'mediapipe_model_maker')
+
+
+def _parse_requirements(path):
+  with open(os.path.join(MM_ROOT_PATH, path)) as f:
+    return [
+        line.rstrip()
+        for line in f
+        if not (line.isspace() or line.startswith('#'))
+    ]
+
+
+def _copy_to_pip_src_dir(file):
+  """Copy a file from bazel-bin to the pip_src dir."""
+  dst = file
+  dst_dir = os.path.dirname(dst)
+  if not os.path.exists(dst_dir):
+    os.makedirs(dst_dir)
+  src_file = os.path.join('../../bazel-bin/mediapipe/model_maker', file)
+  shutil.copyfile(src_file, file)
+
+
+def _setup_build_dir():
+  """Setup the BUILD_DIR directory to build the mediapipe_model_maker package.
+
+  We need to create a new BUILD_DIR directory because any references to the path
+  `mediapipe/model_maker` needs to be renamed to `mediapipe_model_maker` to
+  avoid conflicting with the mediapipe package name.
+  This setup function performs the following actions:
+  1. Copy python source code into BUILD_DIR and rename imports to
+    mediapipe_model_maker
+  2. Download models from GCS into BUILD_DIR
+  """
+  # Copy python source code into BUILD_DIR
+  if os.path.exists(BUILD_DIR):
+    shutil.rmtree(BUILD_DIR)
+  python_files = glob.glob('python/**/*.py', recursive=True)
+  python_files.append('__init__.py')
+  for python_file in python_files:
+    # Exclude test files from pip package
+    if '_test.py' in python_file:
+      continue
+    build_target_file = os.path.join(BUILD_MM_DIR, python_file)
+    with open(python_file, 'r') as file:
+      filedata = file.read()
+    # Rename all mediapipe.model_maker imports to mediapipe_model_maker
+    filedata = filedata.replace('from mediapipe.model_maker',
+                                'from mediapipe_model_maker')
+    os.makedirs(os.path.dirname(build_target_file), exist_ok=True)
+    with open(build_target_file, 'w') as file:
+      file.write(filedata)
+
+  # Use bazel to download GCS model files
+  model_build_files = ['models/gesture_recognizer/BUILD']
+  for model_build_file in model_build_files:
+    build_target_file = os.path.join(BUILD_MM_DIR, model_build_file)
+    os.makedirs(os.path.dirname(build_target_file), exist_ok=True)
+    shutil.copy(model_build_file, build_target_file)
+  external_files = [
+      'models/gesture_recognizer/canned_gesture_classifier.tflite',
+      'models/gesture_recognizer/gesture_embedder.tflite',
+      'models/gesture_recognizer/hand_landmark_full.tflite',
+      'models/gesture_recognizer/palm_detection_full.tflite',
+      'models/gesture_recognizer/gesture_embedder/keras_metadata.pb',
+      'models/gesture_recognizer/gesture_embedder/saved_model.pb',
+      'models/gesture_recognizer/gesture_embedder/variables/variables.data-00000-of-00001',
+      'models/gesture_recognizer/gesture_embedder/variables/variables.index',
+    ]
+  for elem in external_files:
+    external_file = os.path.join(f'{SRC_NAME}/mediapipe_model_maker', elem)
+    sys.stderr.write('downloading file: %s\n' % external_file)
+    fetch_model_command = [
+        'bazel',
+        'build',
+        external_file,
+    ]
+    if subprocess.call(fetch_model_command) != 0:
+      sys.exit(-1)
+    _copy_to_pip_src_dir(external_file)
+
+_setup_build_dir()
+
+setuptools.setup(
+    name='mediapipe-model-maker',
+    version=__version__,
+    url='https://github.com/google/mediapipe/tree/master/mediapipe/model_maker',
+    description='MediaPipe Model Maker is a simple, low-code solution for customizing on-device ML models',
+    author='The MediaPipe Authors',
+    author_email='mediapipe@google.com',
+    long_description='',
+    long_description_content_type='text/markdown',
+    packages=setuptools.find_packages(where=SRC_NAME),
+    package_dir={'': SRC_NAME},
+    install_requires=_parse_requirements('requirements.txt'),
+    include_package_data=True,
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Operating System :: MacOS :: MacOS X',
+        'Operating System :: Microsoft :: Windows',
+        'Operating System :: POSIX :: Linux',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3 :: Only',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+    license='Apache 2.0',
+    keywords=['mediapipe', 'model', 'maker'],
+)