安装 | 简单示例 | 详细示例 | 文档 | 贡献 | 许可证

PyTorchCodec

PyTorchCodec 是一个用于将视频解码为 PyTorch 张量的 Python 库。它旨在快速、易于使用，并与 PyTorch 生态系统良好集成。如果您想使用 PyTorch 在视频上训练机器学习模型，那么 TorchCodec 就是您将这些视频转换为数据的工具。

我们通过以下方式实现这些功能：

• 遵循Python和PyTorch约定的Pythonic API。
• 依赖于FFmpeg进行解码。TorchCodec使用您已安装的FFmpeg版本。FFmpeg是一个成熟的库，在大多数系统上都有广泛的覆盖范围。然而，它并不容易使用。TorchCodec抽象FFmpeg的复杂性，以确保它被正确且高效地使用。
• 以PyTorch张量的形式返回数据，可以直接用于训练模型或输入到PyTorch转换中。

[!NOTE] 注意：TorchCodec仍处于早期开发阶段，某些API可能会在未来版本中更新，而无需弃用周期，具体取决于用户反馈。如果您有任何建议或问题，请通过创建一个问题告诉我们！

使用TorchCodec

以下是对您可以使用TorchCodec做什么的简短总结。要获取更详细的示例，请查看我们的文档！

from torchcodec.decoders import SimpleVideoDecoder

decoder = SimpleVideoDecoder("path/to/video.mp4")

decoder.metadata
# VideoStreamMetadata:
#   num_frames: 250
#   duration_seconds: 10.0
#   bit_rate: 31315.0
#   codec: h264
#   average_fps: 25.0
#   ... (truncated output)

len(decoder)  # == decoder.metadata.num_frames!
# 250
decoder.metadata.average_fps  # Note: instantaneous fps can be higher or lower
# 25.0

# Simple Indexing API
decoder[0]  # uint8 tensor of shape [C, H, W]
decoder[0 : -1 : 20]  # uint8 stacked tensor of shape [N, C, H, W]


# Iterate over frames:
for frame in decoder:
    pass

# Indexing, with PTS and duration info
decoder.get_frame_at(len(decoder) - 1)
# Frame:
#   data (shape): torch.Size([3, 400, 640])
#   pts_seconds: 9.960000038146973
#   duration_seconds: 0.03999999910593033

decoder.get_frames_at(start=10, stop=30, step=5)
# FrameBatch:
#   data (shape): torch.Size([4, 3, 400, 640])
#   pts_seconds: tensor([0.4000, 0.6000, 0.8000, 1.0000])
#   duration_seconds: tensor([0.0400, 0.0400, 0.0400, 0.0400])

# Time-based indexing with PTS and duration info
decoder.get_frame_displayed_at(pts_seconds=2)
# Frame:
#   data (shape): torch.Size([3, 400, 640])
#   pts_seconds: 2.0
#   duration_seconds: 0.03999999910593033

您可以使用以下代码片段使用FFmpeg生成视频并尝试使用TorchCodec

fontfile=/usr/share/fonts/dejavu-sans-mono-fonts/DejaVuSansMono-Bold.ttf
output_video_file=/tmp/output_video.mp4

ffmpeg -f lavfi -i \
    color=size=640x400:duration=10:rate=25:color=blue \
    -vf "drawtext=fontfile=${fontfile}:fontsize=30:fontcolor=white:x=(w-text_w)/2:y=(h-text_h)/2:text='Frame %{frame_num}'" \
    ${output_video_file}

安装TorchCodec

注意：如果您在MacOS上，您需要从源代码构建。以下说明假设您正在Linux上。

首先，按照官方说明安装最新稳定版本的PyTorch。

然后

pip install torchcodec

您还需要在您的系统上安装FFmpeg，TorchCodec的解码能力由您的底层FFmpeg安装决定。安装FFmpeg有多种选择，例如。

conda install ffmpeg
# or
conda install ffmpeg -c conda-forge

您的Linux发行版可能也预装了FFmpeg。TorchCodec支持[4, 7]之间的所有主要FFmpeg版本。

计划中的未来工作

我们正在积极开发以下功能

• 为MacOS发布轮，这样MacOS用户就可以使用pip install torchcodec。目前这仅在Linux上受支持，但MacOS用户可以从源代码构建。
• GPU解码
• 音频解码

如果您有任何功能请求，请通过创建一个问题告诉我们！

贡献

我们欢迎对TorchCodec的贡献！请参阅我们的贡献指南以获取更多详细信息。

许可证

TorchCodec在BSD 3许可证下发布。