3901ad8165
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
198 lines
6.8 KiB
Python
198 lines
6.8 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
|
||
from typing import Any, cast
|
||
|
||
import av
|
||
import numpy as np
|
||
|
||
|
||
class FragmentedMP4Muxer:
|
||
"""Incrementally mux video frames into one fragmented MP4 byte stream."""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
width: int,
|
||
height: int,
|
||
fps: float = 25.0,
|
||
video_codec: str = "h264",
|
||
crf: str = "18",
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> None:
|
||
self._buf = io.BytesIO()
|
||
self._closed = False
|
||
self._container = av.open(
|
||
self._buf,
|
||
mode="w",
|
||
format="mp4",
|
||
options={"movflags": "+frag_every_frame+empty_moov+default_base_moof"},
|
||
)
|
||
|
||
self._stream: av.VideoStream = cast(
|
||
av.VideoStream,
|
||
self._container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)),
|
||
)
|
||
self._stream.width = width
|
||
self._stream.height = height
|
||
self._stream.pix_fmt = "yuv420p"
|
||
|
||
options: dict[str, object] = {"crf": str(crf)}
|
||
if video_codec_options:
|
||
options.update(video_codec_options)
|
||
self._stream.options = options
|
||
|
||
try:
|
||
self._stream.codec_context.max_b_frames = 0
|
||
except AttributeError:
|
||
pass
|
||
|
||
def mux_video_frames(self, video_frames: np.ndarray) -> bytes:
|
||
"""Mux a batch of ``uint8`` RGB frames and return newly written MP4 bytes."""
|
||
if self._closed:
|
||
raise RuntimeError("Cannot mux frames after FragmentedMP4Muxer.close().")
|
||
if video_frames.ndim != 4 or video_frames.shape[-1] != 3:
|
||
raise ValueError("video_frames must have shape (T, H, W, 3).")
|
||
if video_frames.dtype != np.uint8:
|
||
raise ValueError("video_frames must be uint8.")
|
||
if video_frames.shape[1] != self._stream.height or video_frames.shape[2] != self._stream.width:
|
||
raise ValueError("All fragmented MP4 chunks in a session must use the same frame size.")
|
||
|
||
for frame_data in video_frames:
|
||
frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
|
||
for packet in self._stream.encode(frame):
|
||
self._container.mux(packet)
|
||
return self._read_new_bytes()
|
||
|
||
def close(self) -> bytes:
|
||
"""Flush delayed encoder packets, close the container, and return final bytes."""
|
||
if self._closed:
|
||
return b""
|
||
for packet in self._stream.encode():
|
||
self._container.mux(packet)
|
||
self._container.close()
|
||
self._closed = True
|
||
return self._read_new_bytes()
|
||
|
||
def _read_new_bytes(self) -> bytes:
|
||
"""Return newly muxed bytes in the current video container,
|
||
then clear the buffer to prepare for the next chunk."""
|
||
chunk = self._buf.getvalue()
|
||
self._buf.seek(0)
|
||
self._buf.truncate()
|
||
return chunk
|
||
|
||
|
||
def finalize_streaming_video_bytes(
|
||
video_bytes: bytes,
|
||
*,
|
||
input_format: str,
|
||
fps: float = 25.0,
|
||
video_codec_options: dict[str, str] | None = None,
|
||
) -> bytes:
|
||
"""Convert streamed video bytes into a progressive MP4 for local playback."""
|
||
if not video_bytes:
|
||
return video_bytes
|
||
|
||
normalized_format = input_format.lower()
|
||
if normalized_format == "m4s":
|
||
demux_format = "mp4"
|
||
else:
|
||
raise ValueError(f"Unsupported streaming video format: {input_format}")
|
||
|
||
try:
|
||
with cast(Any, av.open(io.BytesIO(video_bytes), format=demux_format)) as container:
|
||
stream = container.streams.video[0]
|
||
frame_arrays = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)]
|
||
except Exception:
|
||
return video_bytes
|
||
|
||
if not frame_arrays:
|
||
return video_bytes
|
||
|
||
frames_u8 = np.ascontiguousarray(np.stack(frame_arrays, axis=0), dtype=np.uint8)
|
||
return mux_video_audio_bytes(
|
||
frames_u8,
|
||
None,
|
||
fps=float(fps),
|
||
video_codec_options=video_codec_options,
|
||
)
|
||
|
||
|
||
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 = cast(av.VideoStream, 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: dict[str, object] = {"crf": str(crf)}
|
||
if video_codec_options:
|
||
options.update(video_codec_options)
|
||
v_stream.options = options
|
||
|
||
a_stream: av.AudioStream | None = None
|
||
samples: np.ndarray | None = None
|
||
layout: str | None = 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 = cast(av.AudioStream, 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:
|
||
if samples is None or layout is None:
|
||
raise ValueError("Audio samples were not prepared for muxing.")
|
||
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()
|