cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Testing MediaPipe Python with GPU support

Browse Source
This commit is contained in:
kinaryml 2023-05-17 10:55:40 +05:30
parent 024f782cd9
commit 45addac249
3 changed files with 96 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mediapipe/tasks/python/core/BUILD
View File

@ -34,6 +34,8 @@ py_library(
], ],
deps = [ deps = [
":optional_dependencies", ":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:base_options_py_pb2",
"//mediapipe/tasks/cc/core/proto:external_file_py_pb2", "//mediapipe/tasks/cc/core/proto:external_file_py_pb2",
], ],

27
mediapipe/tasks/python/core/base_options.py
View File

@ -14,13 +14,18 @@
"""Base options for MediaPipe Task APIs.""" """Base options for MediaPipe Task APIs."""
import dataclasses import dataclasses
import enum
import os import os
from typing import Any, Optional 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 base_options_pb2
from mediapipe.tasks.cc.core.proto import external_file_pb2 from mediapipe.tasks.cc.core.proto import external_file_pb2
from mediapipe.tasks.python.core.optional_dependencies import doc_controls 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 _BaseOptionsProto = base_options_pb2.BaseOptions
_ExternalFileProto = external_file_pb2.ExternalFile _ExternalFileProto = external_file_pb2.ExternalFile
@ -42,10 +47,13 @@ class BaseOptions:
model_asset_path: Path to the model asset file. model_asset_path: Path to the model asset file.
model_asset_buffer: The model asset file contents as bytes. 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_path: Optional[str] = None
model_asset_buffer: Optional[bytes] = 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 @doc_controls.do_not_generate_docs
def to_pb2(self) -> _BaseOptionsProto: def to_pb2(self) -> _BaseOptionsProto:
@ -55,17 +63,16 @@ class BaseOptions:
else: else:
full_path = None full_path = None
if self.delegate == BaseOptions.Delegate.GPU:
acceleration_proto = _AccelerationProto(gpu=_DelegateProto.Gpu())
else:
acceleration_proto = _AccelerationProto(tflite=_DelegateProto.TfLite())
return _BaseOptionsProto( return _BaseOptionsProto(
model_asset=_ExternalFileProto( model_asset=_ExternalFileProto(
file_name=full_path, file_content=self.model_asset_buffer)) file_name=full_path, file_content=self.model_asset_buffer),
acceleration=acceleration_proto
@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)
def __eq__(self, other: Any) -> bool: def __eq__(self, other: Any) -> bool:
"""Checks if this object is equal to the given object. """Checks if this object is equal to the given object.

102
setup.py
View File

@ -245,15 +245,28 @@ class BuildModules(build_ext.build_ext):
sys.stderr.write('downloading file: %s\n' % external_file) sys.stderr.write('downloading file: %s\n' % external_file)
self._download_external_file(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 = [ binary_graphs = [
'face_detection/face_detection_short_range_cpu', 'face_detection/face_detection_short_range_gpu',
'face_detection/face_detection_full_range_cpu', 'face_detection/face_detection_full_range_gpu',
'face_landmark/face_landmark_front_cpu', 'face_landmark/face_landmark_front_gpu',
'hand_landmark/hand_landmark_tracking_cpu', 'hand_landmark/hand_landmark_tracking_gpu',
'holistic_landmark/holistic_landmark_cpu', 'objectron/objectron_cpu', 'holistic_landmark/holistic_landmark_gpu', 'objectron/objectron_gpu',
'pose_landmark/pose_landmark_cpu', 'pose_landmark/pose_landmark_gpu',
'selfie_segmentation/selfie_segmentation_cpu' 'selfie_segmentation/selfie_segmentation_gpu'
] ]
for elem in binary_graphs: for elem in binary_graphs:
binary_graph = os.path.join('mediapipe/modules/', elem) binary_graph = os.path.join('mediapipe/modules/', elem)
sys.stderr.write('generating binarypb: %s\n' % binary_graph) sys.stderr.write('generating binarypb: %s\n' % binary_graph)
@ -271,23 +284,42 @@ class BuildModules(build_ext.build_ext):
sys.exit(-1) sys.exit(-1)
_copy_to_build_lib_dir(self.build_lib, external_file) _copy_to_build_lib_dir(self.build_lib, external_file)
def _generate_binary_graph(self, binary_graph_target): # def _generate_binary_graph(self, binary_graph_target):
"""Generate binary graph for a particular MediaPipe 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 = [ def _generate_binary_graph(self, binary_graph_target):
'bazel', """Generate binary graph for a particular MediaPipe binary graph target."""
'build',
'--compilation_mode=opt', bazel_command = [
'--copt=-DNDEBUG', 'bazel',
'--define=MEDIAPIPE_DISABLE_GPU=1', 'build',
'--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable), '--compilation_mode=opt',
binary_graph_target, '--copt=-DNDEBUG',
] '--copt=-DMESA_EGL_NO_X11_HEADERS',
if not self.link_opencv and not IS_WINDOWS: '--copt=-DEGL_NO_X11',
bazel_command.append('--define=OPENCV=source') '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable),
if subprocess.call(bazel_command) != 0: binary_graph_target,
sys.exit(-1) ]
_copy_to_build_lib_dir(self.build_lib, binary_graph_target + '.binarypb') 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): class GenerateMetadataSchema(build_ext.build_ext):
@ -300,11 +332,21 @@ class GenerateMetadataSchema(build_ext.build_ext):
'object_detector_metadata_schema_py', 'object_detector_metadata_schema_py',
'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_command = [
'bazel', 'bazel',
'build', 'build',
'--compilation_mode=opt', '--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), '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable),
'//mediapipe/tasks/metadata:' + target, '//mediapipe/tasks/metadata:' + target,
] ]
@ -385,12 +427,22 @@ class BuildExtension(build_ext.build_ext):
def _build_binary(self, ext, extra_args=None): def _build_binary(self, ext, extra_args=None):
if not os.path.exists(self.build_temp): if not os.path.exists(self.build_temp):
os.makedirs(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_command = [
'bazel', 'bazel',
'build', 'build',
'--compilation_mode=opt', '--compilation_mode=opt',
'--copt=-DNDEBUG', '--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), '--action_env=PYTHON_BIN_PATH=' + _normalize_path(sys.executable),
str(ext.bazel_target + '.so'), str(ext.bazel_target + '.so'),
] ]