Adds Option, to add an audio stream to write_video and show_video.

mediapy/__init__.py

@@ -1609,6 +1609,14 @@ 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. Can be a path to an audio file or a NumPy array.
+      If 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.
+    audio_sample_rate: Sample rate of the audio in Hz. Required if `audio` is a
+      NumPy array.
+    audio_codec: Audio codec to use (e.g., 'aac'). If None, defaults to 'aac' if
+      audio is provided.
   """
 
   def __init__(
@@ -1627,6 +1635,9 @@ def __init__(
       dtype: _DTypeLike = np.uint8,
       encoded_format: str | None = None,
       sandbox_max_run_time_secs: int | None = None,
+      audio: _NDArray | _Path | None = None,
+      audio_sample_rate: int | None = None,
+      audio_codec: str | None = None,
   ) -> None:
     _check_2d_shape(shape)
     if fps is None and metadata:
@@ -1682,6 +1693,10 @@ 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
+    self._audio_temp_dir: tempfile.TemporaryDirectory[str] | None = None
     if num_rate_specifications == 0 and not ffmpeg_args:
       qp = 20 if math.prod(self.shape) <= 640 * 480 else 28
     self._bitrate_args = (
@@ -1713,6 +1728,63 @@ def __enter__(self) -> 'VideoWriter':
       # 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_codec = self.audio_codec or 'aac'
+        audio_output_args = ['-c:a', audio_codec]
+
+        if isinstance(self.audio, (str, os.PathLike)):
+          audio_path = str(self.audio)
+          audio_input_args = ['-i', audio_path]
+          allowed_input_files.append(audio_path)
+        elif isinstance(self.audio, np.ndarray):
+          if self.audio_sample_rate is None:
+            raise ValueError(
+                'audio_sample_rate must be specified for NumPy array audio.'
+            )
+
+          self._audio_temp_dir = tempfile.TemporaryDirectory()
+          dtype_map = {
+              np.int16: 's16le',
+              np.float32: 'f32le',
+              np.uint8: 'u8',
+              np.int32: 's32le',
+              np.float64: 'f64le',
+          }
+          audio_format = dtype_map.get(self.audio.dtype.type)
+          if not audio_format:
+            raise ValueError(f'Unsupported audio dtype: {self.audio.dtype}')
+
+          tmp_audio_path = pathlib.Path(self._audio_temp_dir.name) / 'audio.raw'
+          tmp_audio_path.write_bytes(self.audio.tobytes())
+
+          audio_path = str(tmp_audio_path)
+
+          if self.audio.ndim == 1:
+            channels = 1
+          elif self.audio.ndim == 2:
+            channels = self.audio.shape[1]
+          else:
+            raise ValueError(f'Unsupported audio shape: {self.audio.shape}')
+
+          audio_input_args = [
+              '-f',
+              audio_format,
+              '-ar',
+              str(self.audio_sample_rate),
+              '-ac',
+              str(channels),
+              '-i',
+              audio_path,
+          ]
+          allowed_input_files.append(audio_path)
+        else:
+          raise ValueError('Unsupported audio type.')
+
       command = (
           [
               '-v',
@@ -1729,7 +1801,10 @@ def __enter__(self) -> 'VideoWriter':
               f'{self.fps}',
               '-i',
               '-',
-              '-an',
+          ]
+          + audio_input_args
+          + audio_output_args
+          + [
               '-vcodec',
               self.codec,
               '-pix_fmt',
@@ -1743,6 +1818,9 @@ def __enter__(self) -> 'VideoWriter':
           command,
           stdin=subprocess.PIPE,
           stderr=subprocess.PIPE,
+          allowed_input_files=allowed_input_files
+          if allowed_input_files
+          else None,
           allowed_output_files=[tmp_name],
           sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
       )
@@ -1814,6 +1892,9 @@ def close(self) -> None:
       self._popen.__exit__(None, None, None)
       self._popen = None
       self._proc = None
+    if self._audio_temp_dir:
+      self._audio_temp_dir.cleanup()
+      self._audio_temp_dir = None
     if self._write_via_local_file:
       # pylint: disable-next=no-member
       self._write_via_local_file.__exit__(None, None, None)
@@ -1976,7 +2057,13 @@ def html_from_compressed_video(
 
 
 def show_video(
-    images: Iterable[_NDArray], *, title: str | None = None, **kwargs: Any
+    images: Iterable[_NDArray],
+    *,
+    title: str | None = None,
+    audio: _NDArray | _Path | None = None,
+    audio_sample_rate: int | None = None,
+    audio_codec: str | None = None,
+    **kwargs: Any,
 ) -> str | None:
   """Displays a video in the IPython notebook and optionally saves it to a file.
 
@@ -1993,12 +2080,24 @@ def show_video(
     images: Iterable of video frames (e.g., a 4D array or a list of 2D or 3D
       arrays).
     title: Optional text shown centered above the video.
+    audio: Optional audio data. Can be a path to an audio file or a NumPy array.
+    audio_sample_rate: Sample rate of the audio in Hz. Required if `audio` is a
+      NumPy array.
+    audio_codec: Audio codec to use (e.g., 'aac'). If None, defaults to 'aac' if
+      audio is provided.
     **kwargs: See `show_videos`.
 
   Returns:
     html string if `return_html` is `True`.
   """
-  return show_videos([images], [title], **kwargs)
+  return show_videos(
+      [images],
+      [title],
+      audio=audio,
+      audio_sample_rate=audio_sample_rate,
+      audio_codec=audio_codec,
+      **kwargs,
+  )
 
 
 def show_videos(
@@ -2016,6 +2115,9 @@ def show_videos(
     ylabel: str = '',
     html_class: str = 'show_videos',
     return_html: bool = False,
+    audio: _NDArray | _Path | None = None,
+    audio_sample_rate: int | None = None,
+    audio_codec: str | None = None,
     **kwargs: Any,
 ) -> str | None:
   """Displays a row of videos in the IPython notebook.
@@ -2047,6 +2149,11 @@ 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.
+    audio: Optional audio data. Can be a path to an audio file or a NumPy array.
+    audio_sample_rate: Sample rate of the audio in Hz. Required if `audio` is a
+      NumPy array.
+    audio_codec: Audio codec to use (e.g., 'aac'). If None, defaults to 'aac' if
+      audio is provided.
     **kwargs: Additional parameters (`border`, `loop`, `autoplay`) for
       `html_from_compressed_video`.
 
@@ -2081,7 +2188,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=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)

mediapy_test.py

@@ -576,6 +576,105 @@ 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)
 
+  def test_video_write_with_audio_array(self):
+    shape = 120, 160
+    num_images = 10
+    fps = 30
+    original_video = media.to_uint8(media.moving_circle(shape, num_images))
+
+    sample_rate = 44100
+    duration = num_images / fps
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    audio_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
+
+    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,
+      )
+
+      self.assertTrue(tmp_path.is_file())
+
+      new_video = media.read_video(tmp_path)
+      self.assertEqual(new_video.metadata.num_images, num_images)
+
+      audio_out_path = pathlib.Path(directory_name) / 'extracted.raw'
+      command = [
+          '-i',
+          str(tmp_path),
+          '-vn',
+          '-acodec',
+          'copy',
+          '-f',
+          's16le',
+          '-y',
+          str(audio_out_path),
+      ]
+      with media._run_ffmpeg(
+          command,
+          allowed_input_files=[str(tmp_path)],
+          allowed_output_files=[str(audio_out_path)],
+      ) as proc:
+        proc.wait()
+        self.assertEqual(proc.returncode, 0)
+
+      self.assertTrue(audio_out_path.is_file())
+      self.assertGreater(audio_out_path.stat().st_size, 0)
+
+  def test_video_write_with_stereo_audio_array(self):
+    shape = 120, 160
+    num_images = 10
+    fps = 30
+    original_video = media.to_uint8(media.moving_circle(shape, num_images))
+
+    sample_rate = 44100
+    duration = num_images / fps
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    audio_data_ch1 = np.sin(2 * np.pi * 440 * t) * 32767
+    audio_data_ch2 = np.sin(2 * np.pi * 880 * t) * 32767
+    audio_data = np.stack((audio_data_ch1, audio_data_ch2), axis=1).astype(
+        np.int16
+    )
+
+    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,
+      )
+
+      self.assertTrue(tmp_path.is_file())
+
+      audio_out_path = pathlib.Path(directory_name) / 'extracted.raw'
+      command = [
+          '-i',
+          str(tmp_path),
+          '-vn',
+          '-acodec',
+          'copy',
+          '-f',
+          's16le',
+          '-y',
+          str(audio_out_path),
+      ]
+      with media._run_ffmpeg(
+          command,
+          allowed_input_files=[str(tmp_path)],
+          allowed_output_files=[str(audio_out_path)],
+      ) as proc:
+        proc.wait()
+        self.assertEqual(proc.returncode, 0)
+
+      self.assertTrue(audio_out_path.is_file())
+      self.assertGreater(audio_out_path.stat().st_size, 0)
+
   def test_video_streaming_write_read_roundtrip(self):
     shape = 62, 744
     num_images = 20
@@ -706,6 +805,21 @@ 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  # moving_circle default
+    duration = len(video) / fps
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    audio_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
+
+    htmls = []
+    with mock.patch('IPython.display.display', htmls.append):
+      media.show_video(video, audio=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>')
+
   def test_show_video_gif(self):
     htmls = []
     with mock.patch('IPython.display.display', htmls.append):