使用 VideoEncoder 编码视频帧

使用 VideoEncoder 编码视频帧#

在本示例中，我们将学习如何使用 torchcodec.encoders.VideoEncoder 类将视频帧编码到文件或原始字节。

准备输入帧#

我们先下载一个示例视频，并将其中部分帧解码为张量，这些张量将作为 VideoEncoder 的输入。关于解码的更多细节可参考项目中的解码示例。

import httpx
from torchcodec.decoders import VideoDecoder
from IPython.display import Video

def play_video(encoded_bytes):
    return Video(
        data=encoded_bytes.numpy().tobytes(),
        embed=True,
        width=640,
        height=360,
        mimetype="video/mp4",
    )

视频来源：https://www.pexels.com/video/adorable-cats-on-the-lawn-4977395/（作者：Altaf Shah）

url = "https://videos.pexels.com/video-files/4977395/4977395-hd_1920_1080_24fps.mp4"
response = httpx.get(url, headers={"User-Agent": ""})
if response.status_code != 200:
    raise RuntimeError(f"下载视频失败。 status_code={response.status_code}.")
raw_video_bytes = response.content

decoder = VideoDecoder(raw_video_bytes)
frames = decoder.get_frames_in_range(0, 60).data
frame_rate = decoder.metadata.average_fps

创建编码器#

我们需要提供一个形状为 (num_frames, num_channels, height, width)、取值范围在 [0, 255] 且 torch.uint8 类型的 4D 张量，以及输入视频的帧率。frame_rate 同时也会用于输出视频。

from torchcodec.encoders import VideoEncoder
print(f"frames.shape={frames.shape}, frames.dtype={frames.dtype}")
print(f"frame_rate={frame_rate} fps")
encoder = VideoEncoder(frames=frames, frame_rate=frame_rate)

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[5], line 1
----> 1 from torchcodec.encoders import VideoEncoder
      2 print(f"frames.shape={frames.shape}, frames.dtype={frames.dtype}")
      3 print(f"frame_rate={frame_rate} fps")

ImportError: cannot import name 'VideoEncoder' from 'torchcodec.encoders' (c:\Users\xinzo\.conda\envs\py313\Lib\site-packages\torchcodec\encoders\__init__.py)

编码到文件、字节或类文件对象#

VideoEncoder 支持编码到文件（to_file）、类文件对象（to_file_like）或原始字节（to_tensor）。下面先编码为原始字节以便展示。

encoded_frames = encoder.to_tensor(format="mp4")
play_video(encoded_frames)

现在我们将编码后的数据再解码回来，验证编码/解码的往返过程。

decoder_verify = VideoDecoder(encoded_frames)
decoded_frames = decoder_verify.get_frames_in_range(0, 60).data
print(f"Re-decoded video: decoded_frames.shape={decoded_frames.shape}")
print(f"Original frames: frames.shape={frames.shape}")

编解码器选择#

默认情况下，编解码器会根据目标文件扩展名（to_file 的 dest 参数）或 format 参数（to_file_like 与 to_tensor）自动选择。例如编码为 MP4 时，默认通常为 H.264。

要使用其他编解码器，可设置 codec 参数（如具体实现 "libx264" 或规范 "h264"）。不同编解码器在质量、文件大小与速度之间有不同权衡。

提示：可通过命令行 ffmpeg -encoders 查看系统可用的编码器。

import tempfile
from pathlib import Path
h264_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
encoder.to_file(h264_output, codec="libx264")
hevc_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
encoder.to_file(hevc_output, codec="hevc")
import subprocess
for output, name in [(h264_output, "h264_output"), (hevc_output, "hevc_output")]:
    result = subprocess.run([
        "ffprobe",
        "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=codec_name",
        "-of", "default=noprint_wrappers=1:nokey=1",
        output,
    ], capture_output=True, text=True)
    print(f"Codec used in {name}: {result.stdout.strip()}")

像素格式#

pixel_format 参数控制输出视频的色度采样（色度子采样），影响质量与文件大小。常见格式包括：

  • yuv420p：4:2:0（标准质量，更小体积，兼容性好）

  • yuv444p：4:4:4（完整色度分辨率，更高质量，更大体积）

多数播放设备与平台支持 yuv420p，因此它是最常见选择。提示：像素格式支持取决于所用编解码器，可用 ffmpeg -h encoder=<codec_name> 查看选项。

yuv420_encoded_frames = encoder.to_tensor(format="mp4", codec="libx264", pixel_format="yuv420p")
play_video(yuv420_encoded_frames)

CRF（恒定码率因子）#

crf 控制视频质量：值越低质量越高。以常用的 libx264 为例，取值范围为 0（无损）到 51（最低质量），17 或 18 常被认为是视觉无损，默认是 23。不同编解码器的 CRF 行为与范围可能不同，且并非所有编解码器都支持 CRF。可用 ffmpeg -h encoder=<codec_name> 查看。

high_quality_output = encoder.to_tensor(format="mp4", codec="libx264", crf=0)
play_video(high_quality_output)

low_quality_output = encoder.to_tensor(format="mp4", codec="libx264", crf=50)
play_video(low_quality_output)

预设#

preset 控制编码速度与压缩率的权衡。更快的预设编码更快但文件更大，更慢的预设编码更慢但压缩更好。以 libx264 为例，常见预设包括 ultrafast（最快）、fast、medium（默认）、slow 与 veryslow（最慢、压缩最好）。更多细节参见 FFmpeg 的 H.264 编码指南。

并非所有编解码器都支持预设选项，可用 ffmpeg -h encoder=<codec_name> 查看。

fast_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
encoder.to_file(fast_output, codec="libx264", preset="ultrafast")
print(f"Size of fast encoded file: {Path(fast_output).stat().st_size} bytes")
slow_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
encoder.to_file(slow_output, codec="libx264", preset="veryslow")
print(f"Size of slow encoded file: {Path(slow_output).stat().st_size} bytes")

额外选项#

extra_options 接受一个字典以传递 FFmpeg 编码器的专有选项，用于控制超出通用参数的编码设置。例如对 libx264，可设置 GOP 大小（g）、最大 B 帧数（max_b_frames）、调优（tune）等。可用 ffmpeg -h encoder=<codec_name> 查看所有可用选项。

custom_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
encoder.to_file(
    custom_output,
    codec="libx264",
    extra_options={
        "g": 50,
        "max_b_frames": 0,
        "tune": "fastdecode",
    }
)
目录