Datasets 文档

加载视频数据

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

加载视频数据

视频支持是实验性的,可能会发生变化。

视频数据集具有 Video 类型的列,其中包含 torchvision 对象。

要处理视频数据集,您需要安装 torchvisionav 包。请查看 安装 指南以了解如何安装它们。

当您加载视频数据集并调用视频列时,视频会解码为 torchvision 视频。

>>> from datasets import load_dataset, Video

>>> dataset = load_dataset("path/to/video/folder", split="train")
>>> dataset[0]["video"]
<torchcodec.decoders._video_decoder.VideoDecoder object at 0x14a61d5a0>

请先使用行索引,然后使用 video 列来索引视频数据集 - dataset[0]["video"] - 以避免在数据集中创建所有视频对象。否则,如果数据集很大,这可能是一个缓慢且耗时的过程。

有关如何加载任何类型数据集的指南,请参阅通用加载指南

读取帧

使用 VideoReader 并通过 next() 直接从视频访问帧。

>>> video = dataset[0]["video"]
>>> first_frame = video.get_frame_at(0)
>>> first_frame.data.shape
(3, 240, 320)
>>> first_frame.pts_seconds  # timestamp
0.0

要一次获取多个帧,您可以调用 .get_frames_in_range(start: int, stop: int, step: int)。这将返回一个帧批次。这是获取长帧列表的有效方法,请参阅 torchcodec 文档 以查看更多有效访问数据的功能。

>>> import torch
>>> frames = video.get_frames_in_range(0, 6, 1)
>>> frames.data.shape
torch.Size([5, 3, 240, 320])

还有一个 .get_frames_played_in_range(start_seconds: float, stop_seconds: float) 函数,用于访问在特定时间范围内播放的所有帧。

>>> frames = video.get_frames_played_in_range(.5, 1.2)
>>> frames.data.shape
torch.Size([42, 3, 240, 320])

本地文件

您可以从视频路径加载数据集。使用 cast_column() 函数接受一列视频文件路径,并使用 Video 特征将其解码为 torchcodec 视频。

>>> from datasets import Dataset, Video

>>> dataset = Dataset.from_dict({"video": ["path/to/video_1", "path/to/video_2", ..., "path/to/video_n"]}).cast_column("video", Video())
>>> dataset[0]["video"]
<torchcodec.decoders._video_decoder.VideoDecoder object at 0x14a61e080>

如果您只想加载视频数据集的底层路径而不解码视频对象,请在 Video 特征中设置 decode=False

>>> dataset = dataset.cast_column("video", Video(decode=False))
>>> dataset[0]["video"]
{'bytes': None,
 'path': 'path/to/video/folder/video0.mp4'}

VideoFolder

您还可以使用 VideoFolder 数据集构建器加载数据集,它不需要编写自定义的数据加载器。这使得 VideoFolder 成为快速为不同视觉任务创建和加载包含数千个视频的视频数据集的理想选择。您的视频数据集结构应如下所示:

folder/train/dog/golden_retriever.mp4
folder/train/dog/german_shepherd.mp4
folder/train/dog/chihuahua.mp4

folder/train/cat/maine_coon.mp4
folder/train/cat/bengal.mp4
folder/train/cat/birman.mp4

如果数据集遵循 VideoFolder 结构,则可以直接使用 load_dataset() 加载它。

>>> from datasets import load_dataset

>>> dataset = load_dataset("username/dataset_name")
>>> # OR locally:
>>> dataset = load_dataset("/path/to/folder")

对于本地数据集,这相当于在 load_dataset() 中手动传递 videofolder,并在 data_dir 中传递目录。

>>> dataset = load_dataset("videofolder", data_dir="/path/to/folder")

然后,您可以将视频作为 torchcodec.decoders._video_decoder.VideoDecoder 对象进行访问。

>>> dataset["train"][0]
{"video": <torchcodec.decoders._video_decoder.VideoDecoder object at 0x14a61e080>, "label": 0}

>>> dataset["train"][-1]
{"video": <torchcodec.decoders._video_decoder.VideoDecoder object at 0x14a61e090>, "label": 1}

要忽略元数据文件中的信息,请在 load_dataset() 中设置 drop_metadata=True

>>> from datasets import load_dataset

>>> dataset = load_dataset("username/dataset_with_metadata", drop_metadata=True)

如果您没有元数据文件,VideoFolder 会自动从目录名称中推断标签名称。如果要删除自动创建的标签,请设置 drop_labels=True。在这种情况下,您的数据集将只包含一个视频列。

>>> from datasets import load_dataset

>>> dataset = load_dataset("username/dataset_without_metadata", drop_labels=True)

最后,filters 参数允许您仅加载数据集的子集,基于标签或元数据的条件。如果元数据是 Parquet 格式,这特别有用,因为该格式支持快速过滤。还建议将此参数与 streaming=True 一起使用,因为默认情况下,数据集在过滤前会完全下载。

>>> filters = [("label", "=", 0)]
>>> dataset = load_dataset("username/dataset_name", streaming=True, filters=filters)

有关创建自己的 VideoFolder 数据集的更多信息,请查看 创建视频数据集 指南。

WebDataset

WebDataset 格式基于 TAR 归档文件的文件夹,适用于大型视频数据集。由于其大小,WebDataset 通常以流式模式加载(使用 streaming=True)。

您可以像这样加载 WebDataset:

>>> from datasets import load_dataset

>>> dataset = load_dataset("webdataset", data_dir="/path/to/folder", streaming=True)

视频解码

默认情况下,当您迭代数据集时,视频会按顺序解码为 torchvision VideoReaders。它会顺序解码视频的元数据,并且在您访问视频帧之前不会读取它们。

然而,使用多线程解码可以显著加快数据集的处理速度。

>>> import os
>>> num_threads = num_threads = min(32, (os.cpu_count() or 1) + 4)
>>> dataset = dataset.decode(num_threads=num_threads)
>>> for example in dataset:  # up to 20 times faster !
...     ...

您可以使用 num_threads 启用多线程。这对于加快远程数据流传输特别有用。但是,对于快速磁盘上的本地数据,它可能比 num_threads=0 慢。

如果您对解码为 torchvision VideoReaders 的视频不感兴趣,而是想访问路径/字节,您可以禁用解码。

>>> dataset = dataset.decode(False)

注意:目前 IterableDataset.decode() 仅适用于流式数据集。

< > 在 GitHub 上更新