From 01d87f0867b48754901dc59169dc9c9b6ed0624e Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Wed, 2 Nov 2022 14:06:14 -0700
Subject: [PATCH] Expose MediaPipe Tasks in MediaPipe PyPI packages.

PiperOrigin-RevId: 485690163
---
 mediapipe/tasks/python/__init__.py        | 11 +++++++++++
 mediapipe/tasks/python/vision/__init__.py | 18 ++++++++++++++++++
 setup.py                                  |  6 ++++--
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/mediapipe/tasks/python/__init__.py b/mediapipe/tasks/python/__init__.py
index 65c1214a..77683914 100644
--- a/mediapipe/tasks/python/__init__.py
+++ b/mediapipe/tasks/python/__init__.py
@@ -11,3 +11,14 @@
 # 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 Tasks API."""
+
+from . import components
+from . import core
+from . import vision
+
+BaseOptions = core.base_options.BaseOptions
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del core
diff --git a/mediapipe/tasks/python/vision/__init__.py b/mediapipe/tasks/python/vision/__init__.py
index 65c1214a..def11317 100644
--- a/mediapipe/tasks/python/vision/__init__.py
+++ b/mediapipe/tasks/python/vision/__init__.py
@@ -11,3 +11,21 @@
 # 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 Tasks Vision API."""
+
+import mediapipe.tasks.python.vision.core
+import mediapipe.tasks.python.vision.image_classifier
+import mediapipe.tasks.python.vision.object_detector
+
+ImageClassifier = image_classifier.ImageClassifier
+ImageClassifierOptions = image_classifier.ImageClassifierOptions
+ObjectDetector = object_detector.ObjectDetector
+ObjectDetectorOptions = object_detector.ObjectDetectorOptions
+RunningMode = core.vision_task_running_mode.VisionTaskRunningMode
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del core
+del image_classifier
+del object_detector
+del mediapipe
diff --git a/setup.py b/setup.py
index cb304ba7..cda53b1c 100644
--- a/setup.py
+++ b/setup.py
@@ -121,8 +121,10 @@ def _add_mp_init_files():
   mp_dir_init_file = open(MP_DIR_INIT_PY, 'a')
   mp_dir_init_file.writelines([
       '\n', 'from mediapipe.python import *\n',
-      'import mediapipe.python.solutions as solutions', '\n\n',
-      '__version__ = \'{}\''.format(__version__), '\n'
+      'import mediapipe.python.solutions as solutions \n',
+      'import mediapipe.tasks.python as tasks\n', '\n\n', 'del framework\n',
+      'del gpu\n', 'del modules\n', 'del python\n', 'del mediapipe\n',
+      'del util\n', '__version__ = \'{}\''.format(__version__), '\n'
   ])
   mp_dir_init_file.close()