From 45addac249b065c5dd0bdf4b348990f9a981b101 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Wed, 17 May 2023 10:55:40 +0530 Subject: [PATCH] Testing MediaPipe Python with GPU support --- mediapipe/tasks/python/core/BUILD | 2 + mediapipe/tasks/python/core/base_options.py | 27 ++++-- setup.py | 102 +++++++++++++++----- 3 files changed, 96 insertions(+), 35 deletions(-) diff --git a/mediapipe/tasks/python/core/BUILD b/mediapipe/tasks/python/core/BUILD index f90826354..c0abefdde 100644 --- a/mediapipe/tasks/python/core/BUILD +++ b/mediapipe/tasks/python/core/BUILD @@ -34,6 +34,8 @@ py_library( ], deps = [ ":optional_dependencies", + "//mediapipe/calculators/tensor/inference_calculator_py_pb2", + "//mediapipe/tasks/cc/core/proto:acceleration_py_pb2", "//mediapipe/tasks/cc/core/proto:base_options_py_pb2", "//mediapipe/tasks/cc/core/proto:external_file_py_pb2", ], diff --git a/mediapipe/tasks/python/core/base_options.py b/mediapipe/tasks/python/core/base_options.py index 90ef045b8..8247f15bd 100644 --- a/mediapipe/tasks/python/core/base_options.py +++ b/mediapipe/tasks/python/core/base_options.py @@ -14,13 +14,18 @@ """Base options for MediaPipe Task APIs.""" import dataclasses +import enum import os from typing import Any, Optional +from mediapipe.calculators.tensor import inference_calculator_pb2 +from mediapipe.tasks.cc.core.proto import acceleration_pb2 from mediapipe.tasks.cc.core.proto import base_options_pb2 from mediapipe.tasks.cc.core.proto import external_file_pb2 from mediapipe.tasks.python.core.optional_dependencies import doc_controls +_DelegateProto = inference_calculator_pb2.InferenceCalculatorOptions.Delegate +_AccelerationProto = acceleration_pb2.Acceleration _BaseOptionsProto = base_options_pb2.BaseOptions _ExternalFileProto = external_file_pb2.ExternalFile @@ -42,10 +47,13 @@ class BaseOptions: model_asset_path: Path to the model asset file. model_asset_buffer: The model asset file contents as bytes. """ + class Delegate(enum.Enum): + CPU = 0 + GPU = 1 model_asset_path: Optional[str] = None model_asset_buffer: Optional[bytes] = None - # TODO: Allow Python API to specify acceleration settings. + delegate: Optional[Delegate] = None @doc_controls.do_not_generate_docs def to_pb2(self) -> _BaseOptionsProto: @@ -55,17 +63,16 @@ class BaseOptions: else: full_path = None + if self.delegate == BaseOptions.Delegate.GPU: + acceleration_proto = _AccelerationProto(gpu=_DelegateProto.Gpu()) + else: + acceleration_proto = _AccelerationProto(tflite=_DelegateProto.TfLite()) + return _BaseOptionsProto( model_asset=_ExternalFileProto( - file_name=full_path, file_content=self.model_asset_buffer)) - - @classmethod - @doc_controls.do_not_generate_docs - def create_from_pb2(cls, pb2_obj: _BaseOptionsProto) -> 'BaseOptions': - """Creates a `BaseOptions` object from the given protobuf object.""" - return BaseOptions( - model_asset_path=pb2_obj.model_asset.file_name, - model_asset_buffer=pb2_obj.model_asset.file_content) + file_name=full_path, file_content=self.model_asset_buffer), + acceleration=acceleration_proto + ) def __eq__(self, other: Any) -> bool: """Checks if this object is equal to the given object. diff --git a/setup.py b/setup.py index 892c6dca7..0712a95b0 100644 --- a/setup.py +++ b/setup.py @@ -245,15 +245,28 @@ class BuildModules(build_ext.build_ext): sys.stderr.write('downloading file: %s\n' % external_file) self._download_external_file(external_file) + # CPU binary graphs + # binary_graphs = [ + # 'face_detection/face_detection_short_range_cpu', + # 'face_detection/face_detection_full_range_cpu', + # 'face_landmark/face_landmark_front_cpu', + # 'hand_landmark/hand_landmark_tracking_cpu', + # 'holistic_landmark/holistic_landmark_cpu', 'objectron/objectron_cpu', + # 'pose_landmark/pose_landmark_cpu', + # 'selfie_segmentation/selfie_segmentation_cpu' + # ] + + # GPU binary graphs binary_graphs = [ - 'face_detection/face_detection_short_range_cpu', - 'face_detection/face_detection_full_range_cpu', - 'face_landmark/face_landmark_front_cpu', - 'hand_landmark/hand_landmark_tracking_cpu', - 'holistic_landmark/holistic_landmark_cpu', 'objectron/objectron_cpu', - 'pose_landmark/pose_landmark_cpu', - 'selfie_segmentation/selfie_segmentation_cpu' + 'face_detection/face_detection_short_range_gpu', + 'face_detection/face_detection_full_range_gpu', + 'face_landmark/face_landmark_front_gpu', + 'hand_landmark/hand_landmark_tracking_gpu', + 'holistic_landmark/holistic_landmark_gpu', 'objectron/objectron_gpu', + 'pose_landmark/pose_landmark_gpu', + 'selfie_segmentation/selfie_segmentation_gpu' ] + for elem in binary_graphs: binary_graph = os.path.join('mediapipe/modules/', elem) sys.stderr.write('generating binarypb: %s\n' % binary_graph) @@ -271,23 +284,42 @@ class BuildModules(build_ext.build_ext): sys.exit(-1) _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.""" + # def _generate_binary_graph(self, binary_graph_target): + # """Generate binary graph for a particular MediaPipe binary graph target.""" + # + # bazel_command = [ + # 'bazel', + # 'build', + # '--compilation_mode=opt', + # '--copt=-DNDEBUG', + # '--define=MEDIAPIPE_DISABLE_GPU=1', + # '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), + # binary_graph_target, + # ] + # if not self.link_opencv and not IS_WINDOWS: + # bazel_command.append('--define=OPENCV=source') + # if subprocess.call(bazel_command) != 0: + # sys.exit(-1) + # _copy_to_build_lib_dir(self.build_lib, binary_graph_target + '.binarypb') - bazel_command = [ - 'bazel', - 'build', - '--compilation_mode=opt', - '--copt=-DNDEBUG', - '--define=MEDIAPIPE_DISABLE_GPU=1', - '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), - binary_graph_target, - ] - if not self.link_opencv and not IS_WINDOWS: - bazel_command.append('--define=OPENCV=source') - if subprocess.call(bazel_command) != 0: - sys.exit(-1) - _copy_to_build_lib_dir(self.build_lib, binary_graph_target + '.binarypb') + def _generate_binary_graph(self, binary_graph_target): + """Generate binary graph for a particular MediaPipe binary graph target.""" + + bazel_command = [ + 'bazel', + 'build', + '--compilation_mode=opt', + '--copt=-DNDEBUG', + '--copt=-DMESA_EGL_NO_X11_HEADERS', + '--copt=-DEGL_NO_X11', + '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), + binary_graph_target, + ] + if not self.link_opencv and not IS_WINDOWS: + bazel_command.append('--define=OPENCV=source') + if subprocess.call(bazel_command) != 0: + sys.exit(-1) + _copy_to_build_lib_dir(self.build_lib, binary_graph_target + '.binarypb') class GenerateMetadataSchema(build_ext.build_ext): @@ -300,11 +332,21 @@ class GenerateMetadataSchema(build_ext.build_ext): 'object_detector_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, + # ] + bazel_command = [ 'bazel', 'build', '--compilation_mode=opt', - '--define=MEDIAPIPE_DISABLE_GPU=1', + '--copt=-DMESA_EGL_NO_X11_HEADERS', + '--copt=-DEGL_NO_X11', '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), '//mediapipe/tasks/metadata:' + target, ] @@ -385,12 +427,22 @@ class BuildExtension(build_ext.build_ext): def _build_binary(self, ext, extra_args=None): if not os.path.exists(self.build_temp): os.makedirs(self.build_temp) + # bazel_command = [ + # 'bazel', + # 'build', + # '--compilation_mode=opt', + # '--copt=-DNDEBUG', + # '--define=MEDIAPIPE_DISABLE_GPU=1', + # '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), + # str(ext.bazel_target + '.so'), + # ] bazel_command = [ 'bazel', 'build', '--compilation_mode=opt', '--copt=-DNDEBUG', - '--define=MEDIAPIPE_DISABLE_GPU=1', + '--copt=-DMESA_EGL_NO_X11_HEADERS', + '--copt=-DEGL_NO_X11', '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), str(ext.bazel_target + '.so'), ]