diff --git a/mediapy/__init__.py b/mediapy/__init__.py index 185b3ce..795ab7a 100644 --- a/mediapy/__init__.py +++ b/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) diff --git a/mediapy_test.py b/mediapy_test.py index 1d2fc31..85361e7 100755 --- a/mediapy_test.py +++ b/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).*') # 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)') + def test_show_video_gif(self): htmls = [] with mock.patch('IPython.display.display', htmls.append):