From b239af8bbd67cabb5a4d199ffb15cb19deb042b8 Mon Sep 17 00:00:00 2001
From: Leonhard Helminger <lhelminger@google.com>
Date: Fri, 25 Sep 2026 15:32:17 -0700
Subject: [PATCH] Add optional audio track support to `write_video` and
 `show_video(s)`.

PiperOrigin-RevId: 988530771
---
 mediapy/__init__.py | 204 +++++++++++++++++++++++++++++++++++++-------
 mediapy_test.py     | 193 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 365 insertions(+), 32 deletions(-)

diff --git a/mediapy/__init__.py b/mediapy/__init__.py
index a1b505e..41400f2 100644
--- a/mediapy/__init__.py
+++ b/mediapy/__init__.py
@@ -191,6 +191,17 @@
 _T = typing.TypeVar('_T')
 _Path = typing.Union[str, 'os.PathLike[str]']
 
+# Raw PCM sample formats (as in `ffmpeg -formats`) for the audio dtypes that
+# `VideoWriter` is able to feed to `ffmpeg`.  Multi-byte formats lack the
+# byte-order suffix ('le' or 'be'), which is appended for the native order.
+_FFMPEG_AUDIO_FORMAT_FROM_DTYPE = {
+    np.uint8: 'u8',
+    np.int16: 's16',
+    np.int32: 's32',
+    np.float32: 'f32',
+    np.float64: 'f64',
+}
+
 _IMAGE_COMPARISON_HTML = """\
 <script
   defer
@@ -732,6 +743,23 @@ def _write_via_local_file(path: _Path) -> Iterator[str]:
         f.write(tmp_path.read_bytes())
 
 
+@contextlib.contextmanager
+def _audio_via_local_file(audio: _NDArray) -> Iterator[str]:
+  """Context to write audio samples to a temporary local file.
+
+  Args:
+    audio: Array of raw PCM audio samples, with shape (N,) for mono or (N, C)
+      for C interleaved channels.
+
+  Yields:
+    The name of a local file containing the raw samples.
+  """
+  with tempfile.TemporaryDirectory() as directory_name:
+    tmp_path = pathlib.Path(directory_name) / 'audio.raw'
+    tmp_path.write_bytes(audio.tobytes())
+    yield str(tmp_path)
+
+
 class set_show_save_dir:  # pylint: disable=invalid-name
   """Save all titled output from `show_*()` calls into files.
 
@@ -1609,6 +1637,22 @@ class VideoWriter(_VideoIO):
       'yuv420p' if all shape dimensions are even, else 'yuv444p'.
     sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
       If None, the default limit is 30 minutes. Unused in open source.
+    audio: Optional audio data as a NumPy array.  It should have shape (N,) for
+      mono or (N, C) for multi-channel audio, where N is the number of samples
+      and C is the number of channels.  The dtype must be one of `np.uint8`,
+      `np.int16`, `np.int32`, `np.float32`, or `np.float64`, in native byte
+      order.  Float samples are nominally in [-1.0, 1.0], and values outside
+      this range typically clip; integer samples span the full range of their
+      type, where `np.uint8` is unsigned with silence at 128.  The two streams
+      are not truncated to a common length: if the audio is shorter than the
+      video, the remainder is silent; if it is longer, the audio continues past
+      the end of the video stream (most players keep showing the last frame).
+    audio_sample_rate: Sample rate of the audio in Hz. Required if `audio` is
+      provided.
+    audio_codec: Audio compression algorithm as defined by "ffmpeg -codecs"
+      (default 'aac').  It must be supported by the video container, e.g., 'aac'
+      for MP4 ('h264' or 'hevc') or 'libopus' for WebM ('vp9').  Ignored if
+      `audio` is None.
   """
 
   def __init__(
@@ -1627,6 +1671,9 @@ def __init__(
       dtype: _DTypeLike = np.uint8,
       encoded_format: str | None = None,
       sandbox_max_run_time_secs: int | None = None,
+      audio: _NDArray | None = None,
+      audio_sample_rate: int | None = None,
+      audio_codec: str = 'aac',
   ) -> None:
     _check_2d_shape(shape)
     if fps is None and metadata:
@@ -1660,6 +1707,19 @@ def __init__(
     dtype = np.dtype(dtype)
     if dtype.type not in (np.uint8, np.uint16):
       raise ValueError(f'Type {dtype} is not np.uint8 or np.uint16.')
+    if audio is not None:
+      if audio_sample_rate is None:
+        raise ValueError('The audio_sample_rate must be set if audio is set.')
+      if audio.ndim not in (1, 2):
+        raise ValueError(f'Audio shape {audio.shape} is not (N,) or (N, C).')
+      if audio.dtype.type not in _FFMPEG_AUDIO_FORMAT_FROM_DTYPE:
+        raise ValueError(f'Audio type {audio.dtype} is unsupported.')
+      if not audio.dtype.isnative:
+        raise ValueError(
+            f'Audio type {audio.dtype} is not in native byte order.'
+        )
+      if codec == 'gif':
+        raise ValueError('Audio is not supported with the gif codec.')
     self.path = pathlib.Path(path)
     self.shape = shape
     all_dimensions_are_even = all(dim % 2 == 0 for dim in shape)
@@ -1682,6 +1742,9 @@ def __init__(
     self.dtype = dtype
     self.encoded_format = encoded_format
     self.sandbox_max_run_time_secs = sandbox_max_run_time_secs
+    self.audio = audio
+    self.audio_sample_rate = audio_sample_rate
+    self.audio_codec = audio_codec
     if num_rate_specifications == 0 and not ffmpeg_args:
       qp = 20 if math.prod(self.shape) <= 640 * 480 else 28
     self._bitrate_args = (
@@ -1699,20 +1762,46 @@ def __init__(
       # video_filter = ('split[s0][s1];[s0]palettegen=stats_mode=single[p];'
       #                 '[s1][p]paletteuse=new=1')
       self.ffmpeg_args = ['-vf', video_filter, '-f', 'gif'] + self.ffmpeg_args
-    self._write_via_local_file: Any = None
-    self._popen: subprocess.Popen[bytes] | None = None
+    self._exit_stack: contextlib.ExitStack | None = None
     self._proc: subprocess.Popen[bytes] | None = None
 
   def __enter__(self) -> 'VideoWriter':
     input_pix_fmt = self._get_pix_fmt(self.dtype, self.input_format)
     try:
-      self._write_via_local_file = _write_via_local_file(self.path)
-      # pylint: disable-next=no-member
-      tmp_name = self._write_via_local_file.__enter__()
+      self._exit_stack = contextlib.ExitStack()
+      tmp_name = self._exit_stack.enter_context(
+          _write_via_local_file(self.path)
+      )
 
       # Writing to stdout using ('-f', 'mp4', '-') would require
       # ('-movflags', 'frag_keyframe+empty_moov') and the result is nonportable.
       height, width = self.shape
+
+      audio_input_args = []
+      audio_output_args = ['-an']
+      allowed_input_files = []
+
+      if self.audio is not None:
+        audio_path = self._exit_stack.enter_context(
+            _audio_via_local_file(self.audio)
+        )
+        channels = 1 if self.audio.ndim == 1 else self.audio.shape[1]
+        audio_format = _FFMPEG_AUDIO_FORMAT_FROM_DTYPE[self.audio.dtype.type]
+        if self.audio.dtype.itemsize > 1:
+          audio_format += {'little': 'le', 'big': 'be'}[sys.byteorder]
+        audio_input_args = [
+            '-f',
+            audio_format,
+            '-ar',
+            str(self.audio_sample_rate),
+            '-ac',
+            str(channels),
+            '-i',
+            audio_path,
+        ]
+        audio_output_args = ['-c:a', self.audio_codec]
+        allowed_input_files.append(audio_path)
+
       command = (
           [
               '-v',
@@ -1729,7 +1818,10 @@ def __enter__(self) -> 'VideoWriter':
               f'{self.fps}',
               '-i',
               '-',
-              '-an',
+          ]
+          + audio_input_args
+          + audio_output_args
+          + [
               '-vcodec',
               self.codec,
               '-pix_fmt',
@@ -1739,14 +1831,18 @@ def __enter__(self) -> 'VideoWriter':
           + self.ffmpeg_args
           + ['-y', tmp_name]
       )
-      self._popen = _run_ffmpeg(
-          command,
-          stdin=subprocess.PIPE,
-          stderr=subprocess.PIPE,
-          allowed_output_files=[tmp_name],
-          sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
+      self._proc = self._exit_stack.enter_context(
+          _run_ffmpeg(
+              command,
+              stdin=subprocess.PIPE,
+              stderr=subprocess.PIPE,
+              # `_run_ffmpeg` omits the sandbox flag only for None, so an empty
+              # list would pass an empty '--sandbox_read_access_files'.
+              allowed_input_files=allowed_input_files or None,
+              allowed_output_files=[tmp_name],
+              sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
+          )
       )
-      self._proc = self._popen.__enter__()
     except Exception:
       self.__exit__(None, None, None)
       raise
@@ -1801,23 +1897,24 @@ def add_image(self, image: _NDArray) -> None:
 
   def close(self) -> None:
     """Finishes writing the video.  (Called automatically at end of context.)"""
-    if self._popen:
-      assert self._proc, 'Error: closing an already closed context.'
-      stdin = self._proc.stdin
-      assert stdin is not None
-      stdin.close()
-      if self._proc.wait():
-        stderr = self._proc.stderr
-        assert stderr is not None
-        s = stderr.read().decode('utf-8')
-        raise RuntimeError(f"Error writing '{self.path}': {s}")
-      self._popen.__exit__(None, None, None)
-      self._popen = None
-      self._proc = None
-    if self._write_via_local_file:
-      # pylint: disable-next=no-member
-      self._write_via_local_file.__exit__(None, None, None)
-      self._write_via_local_file = None
+    if self._exit_stack is None:
+      return
+    # Unwinding the stack terminates the `ffmpeg` process, removes the
+    # temporary audio file, and copies the encoded video to a remote `path`.
+    # Raising within the `with` propagates the error into those contexts, so
+    # an incomplete video is discarded rather than copied.
+    with self._exit_stack:
+      self._exit_stack = None
+      if self._proc is not None:
+        proc, self._proc = self._proc, None
+        stdin = proc.stdin
+        assert stdin is not None
+        stdin.close()
+        if proc.wait():
+          stderr = proc.stderr
+          assert stderr is not None
+          s = stderr.read().decode('utf-8')
+          raise RuntimeError(f"Error writing '{self.path}': {s}")
 
 
 class _VideoArray(np.ndarray):
@@ -2016,6 +2113,9 @@ def show_videos(
     ylabel: str = '',
     html_class: str = 'show_videos',
     return_html: bool = False,
+    audios: Iterable[_NDArray] | Mapping[str, _NDArray] | None = None,
+    audio_sample_rate: int | None = None,
+    audio_codec: str = 'aac',
     **kwargs: Any,
 ) -> str | None:
   """Displays a row of videos in the IPython notebook.
@@ -2047,6 +2147,16 @@ def show_videos(
     ylabel: Text (rotated by 90 degrees) shown on the left of each row.
     html_class: CSS class name used in definition of HTML element.
     return_html: If `True` return the raw HTML `str` instead of displaying.
+    audios: Optional iterable of audio tracks, or dictionary of `{title:
+      audio}`; see `VideoWriter`.  Each track is attached to the corresponding
+      video, so an iterable must have one entry (possibly None) per video.  With
+      a dictionary, its keys must match the video titles exactly (a value may be
+      None for no audio).  Audio is incompatible with `codec` 'gif'.  With the
+      default `autoplay=True`, the video starts muted (browsers block unmuted
+      autoplay); unmute it in the player controls, or pass `autoplay=False`.
+    audio_sample_rate: Sample rate of the audio in Hz, shared by all tracks.
+      Required if `audios` is provided.
+    audio_codec: Audio compression algorithm (default 'aac'); see `VideoWriter`.
     **kwargs: Additional parameters (`border`, `loop`, `autoplay`) for
       `html_from_compressed_video`.
 
@@ -2068,11 +2178,33 @@ def show_videos(
           'Number of videos does not match number of titles'
           f' ({len(list_videos)} vs {len(list_titles)}).'
       )
+
+  if audios is None:
+    list_audios = [None] * len(list_videos)
+  elif isinstance(audios, Mapping):
+    missing = set(list_titles).difference(audios)
+    extra = set(audios).difference(list_titles)
+    if missing or extra:
+      raise ValueError(
+          'The audios dictionary keys must match the video titles (use None as'
+          f' the value for no audio); missing: {sorted(missing, key=str)},'
+          f' extra: {sorted(extra, key=str)}.'
+      )
+    list_audios = [audios.get(title) for title in list_titles]  # pyrefly: ignore[bad-argument-type]
+  else:
+    list_audios = list(audios)
+
+  if len(list_videos) != len(list_audios):
+    raise ValueError(
+        'Number of videos does not match number of audio'
+        f' ({len(list_videos)} vs {len(list_audios)}).'
+    )
+
   if codec not in {'h264', 'gif'}:
     raise ValueError(f'Codec {codec} is neither h264 or gif.')
 
   html_strings = []
-  for video, title in zip(list_videos, list_titles):
+  for video, title, video_audio in zip(list_videos, list_titles, list_audios):
     metadata: VideoMetadata | None = getattr(video, 'metadata', None)
     first_image, video = _peek_first(video)
     w, h = _get_width_height(width, height, first_image.shape[:2])
@@ -2081,7 +2213,15 @@ def show_videos(
       video = [resize_image(image, (h, w)) for image in video]
       first_image = video[0]
     data = compress_video(
-        video, metadata=metadata, fps=fps, bps=bps, qp=qp, codec=codec
+        video,
+        metadata=metadata,
+        fps=fps,
+        bps=bps,
+        qp=qp,
+        codec=codec,
+        audio=video_audio,
+        audio_sample_rate=audio_sample_rate,
+        audio_codec=audio_codec,
     )
     if title is not None and _config.show_save_dir:
       suffix = _filename_suffix_from_codec(codec)
diff --git a/mediapy_test.py b/mediapy_test.py
index f33c897..9a02299 100755
--- a/mediapy_test.py
+++ b/mediapy_test.py
@@ -14,6 +14,7 @@
 
 """Tests for package mediapy."""
 
+import base64
 import io
 import pathlib
 import re
@@ -52,6 +53,56 @@ def _rms_diff(a, b) -> float:
   return np.sqrt(np.mean(np.square(a - b)))
 
 
+def _sine_audio(num_samples, frequency, sample_rate) -> np.ndarray:
+  """Returns int16 samples of a sine wave of `frequency` Hz."""
+  t = np.arange(num_samples) / sample_rate
+  return (np.sin(2 * np.pi * frequency * t) * 32767).astype(np.int16)
+
+
+def _decode_audio(test, video_path, out_path, sample_rate) -> np.ndarray:
+  """Returns the int16 samples of the audio stream decoded from a video.
+
+  Args:
+    test: Test case used to verify that `ffmpeg` succeeds.
+    video_path: Input video file.
+    out_path: Temporary file for the decoded raw samples.
+    sample_rate: Sample rate in Hz to which the stream is resampled, so that an
+      incorrectly stored rate changes the number of returned samples.
+
+  Returns:
+    The samples of all channels, interleaved.
+  """
+  # Not '-acodec copy', which would relabel the AAC bitstream as raw PCM
+  # instead of decoding it.  The channel count is left unspecified so that the
+  # returned size reflects the number of channels in the stream.
+  command = [
+      '-i',
+      str(video_path),
+      '-vn',
+      '-f',
+      's16le',
+      '-ar',
+      str(sample_rate),
+      '-y',
+      str(out_path),
+  ]
+  with media._run_ffmpeg(
+      command,
+      allowed_input_files=[str(video_path)],
+      allowed_output_files=[str(out_path)],
+  ) as proc:
+    proc.wait()
+    test.assertEqual(proc.returncode, 0)
+  return np.frombuffer(out_path.read_bytes(), dtype=np.int16)
+
+
+def _dominant_frequency(samples, sample_rate) -> float:
+  """Returns the frequency in Hz with the largest magnitude in `samples`."""
+  spectrum = np.abs(np.fft.rfft(samples.astype(np.float64)))
+  frequencies = np.fft.rfftfreq(len(samples), 1 / sample_rate)
+  return float(frequencies[spectrum.argmax()])
+
+
 class MediapyTest(parameterized.TestCase):
   """Tests for mediapy package."""
 
@@ -576,6 +627,89 @@ def test_video_non_streaming_write_read_roundtrip(
       self.assertGreater(new_video.metadata.bps, 1_000)  # pyrefly: ignore[no-matching-overload]
       self._check_similar(original_video, new_video, max_rms)
 
+  @parameterized.named_parameters(
+      ('mono', (440.0,)),
+      ('stereo', (440.0, 880.0)),
+  )
+  def test_video_write_with_audio_array(self, frequencies):
+    shape = 120, 160
+    num_images = 10
+    fps = 30
+    sample_rate = 44100
+    num_channels = len(frequencies)
+    num_samples = int(sample_rate * num_images / fps)
+    original_video = media.to_uint8(media.moving_circle(shape, num_images))
+    channels = [_sine_audio(num_samples, f, sample_rate) for f in frequencies]
+    audio_data = channels[0] if num_channels == 1 else np.stack(channels, -1)
+
+    with tempfile.TemporaryDirectory() as directory_name:
+      tmp_path = pathlib.Path(directory_name) / 'test.mp4'
+      media.write_video(
+          tmp_path,
+          original_video,
+          fps=fps,
+          audio=audio_data,
+          audio_sample_rate=sample_rate,
+      )
+      new_video = media.read_video(tmp_path)
+      assert new_video.metadata
+      self.assertEqual(new_video.metadata.num_images, num_images)
+
+      raw_path = pathlib.Path(directory_name) / 'audio.raw'
+      decoded = _decode_audio(self, tmp_path, raw_path, sample_rate)
+
+      # The number of samples is only approximate because the AAC encoder pads
+      # the signal, but it is far from `num_samples * other_num_channels`, so
+      # it detects both a dropped channel and an incorrect sample rate.
+      self.assertBetween(
+          decoded.size,
+          (num_samples - 4096) * num_channels,
+          (num_samples + 4096) * num_channels,
+      )
+      self.assertEqual(decoded.size % num_channels, 0)
+      decoded = decoded.reshape(-1, num_channels)
+      for channel, frequency in enumerate(frequencies):
+        self.assertAlmostEqual(
+            _dominant_frequency(decoded[:, channel], sample_rate),
+            frequency,
+            delta=5.0,
+        )
+
+  def test_video_write_with_audio_codec(self):
+    shape = 64, 64
+    num_images = 10
+    fps = 30
+    sample_rate = 48000  # A sample rate natively supported by Opus.
+    num_samples = int(sample_rate * num_images / fps)
+    video = media.to_uint8(media.moving_circle(shape, num_images))
+    audio_data = _sine_audio(num_samples, 440.0, sample_rate)
+
+    with tempfile.TemporaryDirectory() as directory_name:
+      # The default 'aac' audio codec is not supported in a WebM container.
+      tmp_path = pathlib.Path(directory_name) / 'test.webm'
+      media.write_video(
+          tmp_path,
+          video,
+          fps=fps,
+          codec='vp9',
+          audio=audio_data,
+          audio_sample_rate=sample_rate,
+          audio_codec='libopus',
+      )
+      raw_path = pathlib.Path(directory_name) / 'audio.raw'
+      decoded = _decode_audio(self, tmp_path, raw_path, sample_rate)
+      self.assertAlmostEqual(
+          _dominant_frequency(decoded, sample_rate), 440.0, delta=5.0
+      )
+
+  def test_video_writer_audio_non_native_byte_order(self):
+    non_native_dtype = np.dtype(np.int16).newbyteorder()
+    audio_data = np.zeros(1000, dtype=non_native_dtype)
+    with self.assertRaisesRegex(ValueError, 'native byte order'):
+      media.VideoWriter(
+          'test.mp4', (16, 16), audio=audio_data, audio_sample_rate=44100
+      )
+
   def test_video_streaming_write_read_roundtrip(self):
     shape = 62, 744
     num_images = 20
@@ -706,6 +840,22 @@ def test_show_video(self):
     self.assertLen(re.findall('(?s)<video', htmls[0].data), 1)  # pyrefly: ignore[no-matching-overload]
     self.assertRegex(htmls[0].data, '(?s)<video .*>.*</video>')  # pyrefly: ignore[bad-specialization]
 
+  def test_show_video_with_audio(self):
+    video = media.moving_circle()
+    sample_rate = 44100
+    fps = 30
+    num_samples = int(sample_rate * len(video) / fps)
+    audio_data = _sine_audio(num_samples, 440.0, sample_rate)
+
+    htmls = []
+    with mock.patch('IPython.display.display', htmls.append):
+      media.show_video(
+          video, fps=fps, audios=[audio_data], audio_sample_rate=sample_rate
+      )
+    self.assertLen(htmls, 1)
+    self.assertIsInstance(htmls[0], IPython.display.HTML)
+    self.assertRegex(htmls[0].data, '(?s)<video .*>.*</video>')  # pyrefly: ignore[bad-specialization]
+
   def test_show_video_gif(self):
     htmls = []
     with mock.patch('IPython.display.display', htmls.append):
@@ -758,6 +908,49 @@ def test_show_videos_dict(self):
     self.assertRegex(htmls[0].data, '(?s)title1.*<video.*title2.*<video')  # pyrefly: ignore[bad-specialization]
     self.assertLen(re.findall('(?s)<video', htmls[0].data), 2)  # pyrefly: ignore[no-matching-overload]
 
+  @parameterized.named_parameters(('list', False), ('dict', True))
+  def test_show_videos_with_audio(self, use_dict):
+    video = media.moving_circle((32, 32), num_images=10)
+    fps = 30
+    sample_rate = 44100
+    num_samples = int(sample_rate * len(video) / fps)
+    frequencies = [440.0, 880.0]
+    audio = [_sine_audio(num_samples, f, sample_rate) for f in frequencies]
+    titles = ['title1', 'title2']
+    if use_dict:
+      videos = dict(zip(titles, [video] * 2))
+      audios = dict(zip(titles, audio))
+    else:
+      videos = [video] * 2
+      audios = audio
+
+    htmls = []
+    with mock.patch('IPython.display.display', htmls.append):
+      media.show_videos(
+          videos, fps=fps, audios=audios, audio_sample_rate=sample_rate
+      )
+    self.assertLen(htmls, 1)
+    self.assertIsInstance(htmls[0], IPython.display.HTML)
+    self.assertLen(re.findall('(?s)<video', htmls[0].data), 2)  # pyrefly: ignore[no-matching-overload]
+
+    # Decode the embedded videos to verify that each one carries its own track
+    # rather than a single shared one.
+    b64_videos = re.findall('base64,([^"]+)"', htmls[0].data)  # pyrefly: ignore[no-matching-overload]
+    self.assertLen(b64_videos, 2)
+    with tempfile.TemporaryDirectory() as directory_name:
+      for index, (b64_video, frequency) in enumerate(
+          zip(b64_videos, frequencies)
+      ):
+        tmp_path = pathlib.Path(directory_name) / f'video{index}.mp4'
+        tmp_path.write_bytes(base64.b64decode(b64_video))
+        raw_path = pathlib.Path(directory_name) / f'audio{index}.raw'
+        decoded = _decode_audio(self, tmp_path, raw_path, sample_rate)
+        # Approximate because the AAC encoder pads the signal.
+        self.assertBetween(decoded.size, num_samples - 4096, num_samples + 4096)
+        self.assertAlmostEqual(
+            _dominant_frequency(decoded, sample_rate), frequency, delta=5.0
+        )
+
   def test_show_videos_over_multiple_rows(self):
     htmls = []
     with mock.patch('IPython.display.display', htmls.append):