From 57759aedafb19c1c1703968d616dfcc26460ca22 Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Fri, 4 Nov 2022 09:45:33 -0700
Subject: [PATCH] Enable using tasks metadata components in MediaPipe PyPI
 packages.

PiperOrigin-RevId: 486162986
---
 requirements.txt |  1 +
 setup.py         | 51 +++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 00e51ffd6..85a08aea9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 absl-py
 attrs>=19.1.0
+flatbuffers>=2.0
 matplotlib
 numpy
 opencv-contrib-python
diff --git a/setup.py b/setup.py
index cda53b1c6..40c4a7361 100644
--- a/setup.py
+++ b/setup.py
@@ -129,6 +129,15 @@ def _add_mp_init_files():
   mp_dir_init_file.close()
 
 
+def _copy_to_build_lib_dir(build_lib, file):
+  """Copy a file from bazel-bin to the build lib dir."""
+  dst = os.path.join(build_lib + '/', file)
+  dst_dir = os.path.dirname(dst)
+  if not os.path.exists(dst_dir):
+    os.makedirs(dst_dir)
+  shutil.copyfile(os.path.join('bazel-bin/', file), dst)
+
+
 class GeneratePyProtos(build_ext.build_ext):
   """Generate MediaPipe Python protobuf files by Protocol Compiler."""
 
@@ -259,7 +268,7 @@ class BuildModules(build_ext.build_ext):
     ]
     if subprocess.call(fetch_model_command) != 0:
       sys.exit(-1)
-    self._copy_to_build_lib_dir(external_file)
+    _copy_to_build_lib_dir(self.build_lib, external_file)
 
   def _generate_binary_graph(self, binary_graph_target):
     """Generate binary graph for a particular MediaPipe binary graph target."""
@@ -277,15 +286,27 @@ class BuildModules(build_ext.build_ext):
       bazel_command.append('--define=OPENCV=source')
     if subprocess.call(bazel_command) != 0:
       sys.exit(-1)
-    self._copy_to_build_lib_dir(binary_graph_target + '.binarypb')
+    _copy_to_build_lib_dir(self.build_lib, binary_graph_target + '.binarypb')
 
-  def _copy_to_build_lib_dir(self, file):
-    """Copy a file from bazel-bin to the build lib dir."""
-    dst = os.path.join(self.build_lib + '/', file)
-    dst_dir = os.path.dirname(dst)
-    if not os.path.exists(dst_dir):
-      os.makedirs(dst_dir)
-    shutil.copyfile(os.path.join('bazel-bin/', file), dst)
+
+class GenerateMetadataSchema(build_ext.build_ext):
+  """Generate metadata python schema files."""
+
+  def run(self):
+    for target in ['metadata_schema_py', 'schema_py']:
+      bazel_command = [
+          'bazel',
+          'build',
+          '--compilation_mode=opt',
+          '--define=MEDIAPIPE_DISABLE_GPU=1',
+          '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable),
+          '//mediapipe/tasks/metadata:' + target,
+      ]
+      if subprocess.call(bazel_command) != 0:
+        sys.exit(-1)
+      _copy_to_build_lib_dir(
+          self.build_lib,
+          'mediapipe/tasks/metadata/' + target + '_generated.py')
 
 
 class BazelExtension(setuptools.Extension):
@@ -375,6 +396,7 @@ class BuildPy(build_py.build_py):
     build_ext_obj = self.distribution.get_command_obj('build_ext')
     build_ext_obj.link_opencv = self.link_opencv
     self.run_command('gen_protos')
+    self.run_command('generate_metadata_schema')
     self.run_command('build_modules')
     self.run_command('build_ext')
     build_py.build_py.run(self)
@@ -434,18 +456,25 @@ setuptools.setup(
     author_email='mediapipe@google.com',
     long_description=_get_long_description(),
     long_description_content_type='text/markdown',
-    packages=setuptools.find_packages(exclude=['mediapipe.examples.desktop.*']),
+    packages=setuptools.find_packages(
+        exclude=['mediapipe.examples.desktop.*', 'mediapipe.model_maker.*']),
     install_requires=_parse_requirements('requirements.txt'),
     cmdclass={
         'build_py': BuildPy,
-        'gen_protos': GeneratePyProtos,
         'build_modules': BuildModules,
         'build_ext': BuildExtension,
+        'generate_metadata_schema': GenerateMetadataSchema,
+        'gen_protos': GeneratePyProtos,
         'install': Install,
         'restore': Restore,
     },
     ext_modules=[
         BazelExtension('//mediapipe/python:_framework_bindings'),
+        BazelExtension(
+            '//mediapipe/tasks/cc/metadata/python:_pywrap_metadata_version'),
+        BazelExtension(
+            '//mediapipe/tasks/python/metadata/flatbuffers_lib:_pywrap_flatbuffers'
+        ),
     ],
     zip_safe=False,
     include_package_data=True,