81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
from fractions import Fraction
|
||
|
||
import av
|
||
import numpy as np
|
||
|
||
|
||
def mux_video_audio_bytes(
|
||
video_frames: np.ndarray,
|
||
audio_waveform: np.ndarray | None = None,
|
||
*,
|
||
fps: float = 25.0,
|
||
audio_sample_rate: int = 44100,
|
||
video_codec: str = "h264",
|
||
audio_codec: str = "aac",
|
||
crf: str = "18",
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> bytes:
|
||
"""Mux video frames and optional audio waveform into MP4 bytes.
|
||
|
||
Args:
|
||
video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB).
|
||
audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``.
|
||
fps: Video frame rate.
|
||
audio_sample_rate: Audio sample rate in Hz.
|
||
video_codec: Video codec name.
|
||
audio_codec: Audio codec name.
|
||
crf: Constant rate factor for the video encoder.
|
||
|
||
Returns:
|
||
Raw MP4 bytes ready to be written to disk or streamed.
|
||
"""
|
||
buf = io.BytesIO()
|
||
container = av.open(buf, mode="w", format="mp4")
|
||
|
||
v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))
|
||
v_stream.width = video_frames.shape[2]
|
||
v_stream.height = video_frames.shape[1]
|
||
v_stream.pix_fmt = "yuv420p"
|
||
|
||
options = {"crf": str(crf)}
|
||
if video_codec_options:
|
||
options.update(video_codec_options)
|
||
v_stream.options = options
|
||
|
||
a_stream = None
|
||
if audio_waveform is not None:
|
||
samples = audio_waveform.astype(np.float32)
|
||
if samples.ndim == 1:
|
||
samples = samples.reshape(1, -1)
|
||
elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
|
||
samples = np.ascontiguousarray(samples.T)
|
||
num_channels = samples.shape[0]
|
||
layout = "stereo" if num_channels >= 2 else "mono"
|
||
a_stream = container.add_stream(audio_codec, rate=audio_sample_rate)
|
||
a_stream.layout = layout
|
||
|
||
for frame_data in video_frames:
|
||
frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
|
||
for packet in v_stream.encode(frame):
|
||
container.mux(packet)
|
||
for packet in v_stream.encode():
|
||
container.mux(packet)
|
||
|
||
if a_stream is not None and audio_waveform is not None:
|
||
audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
|
||
audio_frame.sample_rate = audio_sample_rate
|
||
for packet in a_stream.encode(audio_frame):
|
||
container.mux(packet)
|
||
for packet in a_stream.encode():
|
||
container.mux(packet)
|
||
|
||
container.close()
|
||
return buf.getvalue()
|