Change swig dependency in pip_packge build to use pybind11.

PiperOrigin-RevId: 302850381
Change-Id: If7693293985c24d0e57c0129d1a73bc6e4ddc4a9

tensorflow/lite/tools/pip_package/Dockerfile

@@ -14,6 +14,7 @@ RUN apt-get update && \
       python-setuptools \
       python-wheel \
       python-numpy \
+      python-pip \
       libpython-dev \
       libpython-dev:armhf \
       libpython-dev:arm64 \
@@ -21,6 +22,7 @@ RUN apt-get update && \
       python3-setuptools \
       python3-wheel \
       python3-numpy \
+      python3-pip \
       libpython3-dev \
       libpython3-dev:armhf \
       libpython3-dev:arm64 \
@@ -29,8 +31,11 @@ RUN apt-get update && \
       zlib1g-dev  \
       zlib1g-dev:armhf \
       zlib1g-dev:arm64 \
-      swig \
       curl \
       unzip \
       git && \
     apt-get clean
+RUN pip install pip --upgrade
+RUN pip install pybind11
+RUN pip3 install pip --upgrade
+RUN pip3 install pybind11

tensorflow/lite/tools/pip_package/setup.py

@@ -24,9 +24,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import glob
 import multiprocessing
 import os
 import subprocess
+import sys
 
 from distutils.command.build_ext import build_ext
 import numpy
@@ -73,9 +75,9 @@ def get_build_cpus():
 
 def make_args(target='', quiet=True):
   """Construct make command line."""
-  args = (['make', 'SHELL=/bin/bash',
+  args = ([
-           'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR]
+      'make', 'SHELL=/bin/bash', 'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR
-          + MAKE_CROSS_OPTIONS +
+  ] + MAKE_CROSS_OPTIONS +
           ['-f', RELATIVE_MAKEFILE_PATH, '-j',
            str(get_build_cpus())])
   if quiet:
@@ -128,28 +130,46 @@ class CustomBuildPy(build_py, object):
     return super(CustomBuildPy, self).run()
 
 
+def get_pybind_include():
+  """pybind11 include directory is not correctly resolved.
+
+  This fixes include directory to /usr/local/pythonX.X
+
+  Returns:
+    include directories to find pybind11
+  """
+  if sys.version_info[0] == 3:
+    include_dirs = glob.glob('/usr/local/include/python3*')
+  else:
+    include_dirs = glob.glob('/usr/local/include/python2*')
+  for include_dir in include_dirs:
+    os.symlink(include_dir, os.path.join(include_dir, 'include'))
+
+  return include_dirs
+
+
 LIB_TFLITE = 'tensorflow-lite'
 LIB_TFLITE_DIR = make_output('libdir')
 
 ext = Extension(
     name='%s._interpreter_wrapper' % PACKAGE_NAME,
     language='c++',
-    sources=['interpreter_wrapper/interpreter_wrapper.i',
+    sources=[
         'interpreter_wrapper/interpreter_wrapper.cc',
+        'interpreter_wrapper/interpreter_wrapper_pybind11.cc',
         'interpreter_wrapper/numpy.cc',
         'interpreter_wrapper/python_error_reporter.cc',
-             'interpreter_wrapper/python_utils.cc'],
+        'interpreter_wrapper/python_utils.cc'
+    ],
     extra_compile_args=['--std=c++11'],
-    swig_opts=['-c++',
+    include_dirs=[
-               '-I%s' % TENSORFLOW_DIR,
+        TENSORFLOW_DIR,
-               '-module', 'interpreter_wrapper',
-               '-outdir', PACKAGE_NAME],
-    include_dirs=[TENSORFLOW_DIR,
         os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
                      'pip_package'),
         numpy.get_include(),
         os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
-                  os.path.join(DOWNLOADS_DIR, 'absl')],
+        os.path.join(DOWNLOADS_DIR, 'absl')
+    ] + get_pybind_include(),
     libraries=[LIB_TFLITE],
     library_dirs=[LIB_TFLITE_DIR])
 
@@ -186,9 +206,9 @@ setup(
     ext_modules=[ext],
     install_requires=[
         'numpy >= 1.16.0',
+        'pybind11 >= 2.4.3',
     ],
     cmdclass={
         'build_ext': CustomBuildExt,
         'build_py': CustomBuildPy,
-    }
+    })
-)