From e4899c820707a3783b82389b22b6ad67e4946cb2 Mon Sep 17 00:00:00 2001 From: Jiuqiang Tang Date: Tue, 1 Nov 2022 08:52:11 -0700 Subject: [PATCH] MediaPipe Python audio base layer. PiperOrigin-RevId: 485329470 --- mediapipe/tasks/python/audio/core/BUILD | 37 ++++++ mediapipe/tasks/python/audio/core/__init__.py | 16 +++ .../audio/core/audio_task_running_mode.py | 29 +++++ .../python/audio/core/base_audio_task_api.py | 123 ++++++++++++++++++ 4 files changed, 205 insertions(+) create mode 100644 mediapipe/tasks/python/audio/core/BUILD create mode 100644 mediapipe/tasks/python/audio/core/__init__.py create mode 100644 mediapipe/tasks/python/audio/core/audio_task_running_mode.py create mode 100644 mediapipe/tasks/python/audio/core/base_audio_task_api.py diff --git a/mediapipe/tasks/python/audio/core/BUILD b/mediapipe/tasks/python/audio/core/BUILD new file mode 100644 index 00000000..d9f169c6 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/BUILD @@ -0,0 +1,37 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +# Placeholder for internal Python strict library and test compatibility macro. + +package(default_visibility = ["//mediapipe/tasks:internal"]) + +licenses(["notice"]) + +py_library( + name = "audio_task_running_mode", + srcs = ["audio_task_running_mode.py"], +) + +py_library( + name = "base_audio_task_api", + srcs = [ + "base_audio_task_api.py", + ], + deps = [ + ":audio_task_running_mode", + "//mediapipe/framework:calculator_py_pb2", + "//mediapipe/python:_framework_bindings", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) diff --git a/mediapipe/tasks/python/audio/core/__init__.py b/mediapipe/tasks/python/audio/core/__init__.py new file mode 100644 index 00000000..6a840518 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/__init__.py @@ -0,0 +1,16 @@ +"""Copyright 2022 The MediaPipe Authors. + +All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +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. +""" diff --git a/mediapipe/tasks/python/audio/core/audio_task_running_mode.py b/mediapipe/tasks/python/audio/core/audio_task_running_mode.py new file mode 100644 index 00000000..0fa36d40 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/audio_task_running_mode.py @@ -0,0 +1,29 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""The running mode of MediaPipe Audio Tasks.""" + +import enum + + +class AudioTaskRunningMode(enum.Enum): + """MediaPipe audio task running mode. + + Attributes: + AUDIO_CLIPS: The mode for running a mediapipe audio task on independent + audio clips. + AUDIO_STREAM: The mode for running a mediapipe audio task on an audio + stream, such as from microphone. + """ + AUDIO_CLIPS = 'AUDIO_CLIPS' + AUDIO_STREAM = 'AUDIO_STREAM' diff --git a/mediapipe/tasks/python/audio/core/base_audio_task_api.py b/mediapipe/tasks/python/audio/core/base_audio_task_api.py new file mode 100644 index 00000000..b6a2e0e4 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/base_audio_task_api.py @@ -0,0 +1,123 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 audio task base api.""" + +from typing import Callable, Mapping, Optional + +from mediapipe.framework import calculator_pb2 +from mediapipe.python._framework_bindings import packet as packet_module +from mediapipe.python._framework_bindings import task_runner as task_runner_module +from mediapipe.tasks.python.audio.core import audio_task_running_mode as running_mode_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_TaskRunner = task_runner_module.TaskRunner +_Packet = packet_module.Packet +_RunningMode = running_mode_module.AudioTaskRunningMode + + +class BaseAudioTaskApi(object): + """The base class of the user-facing mediapipe audio task api classes.""" + + def __init__( + self, + graph_config: calculator_pb2.CalculatorGraphConfig, + running_mode: _RunningMode, + packet_callback: Optional[Callable[[Mapping[str, packet_module.Packet]], + None]] = None + ) -> None: + """Initializes the `BaseAudioTaskApi` object. + + Args: + graph_config: The mediapipe audio task graph config proto. + running_mode: The running mode of the mediapipe audio task. + packet_callback: The optional packet callback for getting results + asynchronously in the audio stream mode. + + Raises: + ValueError: The packet callback is not properly set based on the task's + running mode. + """ + if running_mode == _RunningMode.AUDIO_STREAM: + if packet_callback is None: + raise ValueError( + 'The audio task is in audio stream mode, a user-defined result ' + 'callback must be provided.') + elif packet_callback: + raise ValueError( + 'The audio task is in audio clips mode, a user-defined result ' + 'callback should not be provided.') + self._runner = _TaskRunner.create(graph_config, packet_callback) + self._running_mode = running_mode + + def _process_audio_clip( + self, inputs: Mapping[str, _Packet]) -> Mapping[str, _Packet]: + """A synchronous method to process independent audio clips. + + The call blocks the current thread until a failure status or a successful + result is returned. + + Args: + inputs: A dict contains (input stream name, data packet) pairs. + + Returns: + A dict contains (output stream name, data packet) pairs. + + Raises: + ValueError: If the task's running mode is not set to audio clips mode. + """ + if self._running_mode != _RunningMode.AUDIO_CLIPS: + raise ValueError( + 'Task is not initialized with the audio clips mode. Current running mode:' + + self._running_mode.name) + return self._runner.process(inputs) + + def _send_audio_stream_data(self, inputs: Mapping[str, _Packet]) -> None: + """An asynchronous method to send audio stream data to the runner. + + The results will be available in the user-defined results callback. + + Args: + inputs: A dict contains (input stream name, data packet) pairs. + + Raises: + ValueError: If the task's running mode is not set to the audio stream + mode. + """ + if self._running_mode != _RunningMode.AUDIO_STREAM: + raise ValueError( + 'Task is not initialized with the audio stream mode. Current running mode:' + + self._running_mode.name) + self._runner.send(inputs) + + def close(self) -> None: + """Shuts down the mediapipe audio task instance. + + Raises: + RuntimeError: If the mediapipe audio task failed to close. + """ + self._runner.close() + + @doc_controls.do_not_generate_docs + def __enter__(self): + """Return `self` upon entering the runtime context.""" + return self + + @doc_controls.do_not_generate_docs + def __exit__(self, unused_exc_type, unused_exc_value, unused_traceback): + """Shuts down the mediapipe audio task instance on exit of the context manager. + + Raises: + RuntimeError: If the mediapipe audio task failed to close. + """ + self.close()