MediaPipe Python audio base layer.

PiperOrigin-RevId: 485329470

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",
+    ],
+)

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.
+"""

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'

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()