Datasets 文档

加载图像数据

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

加载图像数据

图像数据集具有 Image 类型的列,其中包含 PIL 对象。

要处理图像数据集,您需要安装 vision 依赖项。请参阅 安装指南以了解如何安装。

当您加载图像数据集并调用图像列时,图像将被解码为 PIL 图像

>>> from datasets import load_dataset, Image

>>> dataset = load_dataset("AI-Lab-Makerere/beans", split="train")
>>> dataset[0]["image"]

请先使用行索引,然后再使用 image 列对图像数据集进行索引 —— dataset[0]["image"] —— 以避免对数据集中的所有图像对象进行解码和重采样。否则,如果您拥有大型数据集,这个过程可能会非常缓慢且耗时。

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

本地文件

您可以从图像路径加载数据集。使用 cast_column() 函数来接收图像文件路径列,并使用 Image 特征将其解码为 PIL 图像

>>> from datasets import Dataset, Image

>>> dataset = Dataset.from_dict({"image": ["path/to/image_1", "path/to/image_2", ..., "path/to/image_n"]}).cast_column("image", Image())
>>> dataset[0]["image"]
<PIL.PngImagePlugin.PngImageFile image mode=RGBA size=1200x215 at 0x15E6D7160>]

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

>>> dataset = load_dataset("AI-Lab-Makerere/beans", split="train").cast_column("image", Image(decode=False))
>>> dataset[0]["image"]
{'bytes': None,
 'path': '/root/.cache/huggingface/datasets/downloads/extracted/b0a21163f78769a2cf11f58dfc767fb458fc7cea5c05dccc0144a2c0f0bc1292/train/bean_rust/bean_rust_train.29.jpg'}

ImageFolder

您还可以使用 ImageFolder 数据集构建器加载数据集,这样就不需要编写自定义数据加载器。这使得 ImageFolder 非常适合为不同的视觉任务快速创建和加载包含数千张图像的图像数据集。您的图像数据集结构应该如下所示

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

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

或者,它应该包含元数据,例如

folder/train/metadata.csv
folder/train/0001.png
folder/train/0002.png
folder/train/0003.png

如果数据集遵循 ImageFolder 结构,那么您可以使用 load_dataset() 直接加载它

>>> from datasets import load_dataset

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

对于本地数据集,这等同于在 load_dataset() 中手动传递 imagefolder 以及在 data_dir 中传递目录

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

然后您可以将视频访问为 PIL.Image 对象

>>> dataset["train"][0]
{"image": <PIL.PngImagePlugin.PngImageFile image mode=RGBA size=1200x215 at 0x15E6D7160>, "label": 0}

>>> dataset["train"][-1]
{"image": <PIL.PngImagePlugin.PngImageFile image mode=RGBA size=1200x215 at 0x15E8DAD30>, "label": 1}

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

>>> from datasets import load_dataset

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

如果您没有元数据文件,ImageFolder 会自动从目录名称推断标签名称。如果您想丢弃自动创建的标签,请设置 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)

有关创建您自己的 ImageFolder 数据集的更多信息,请参阅 创建图像数据集指南。

WebDataset

WebDataset 格式基于 TAR 存档文件夹,适用于大型图像数据集。由于其体积较大,WebDataset 通常在流模式下加载(使用 streaming=True)。

您可以像这样加载 WebDataset:

>>> from datasets import load_dataset

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

Lance

Lance 是一种开放的多模态湖仓表格式。Lance 表不仅可以原生存储文本和标量值,还可以将图像、音频和视频等大型二进制对象 (blobs) 与您的表格数据一起存储。

Lance 将您的元数据和图像 blob 存储在一起,同时仍允许您高效地仅扫描您关心的元数据列,而无需加载图像字节。准备好后,您可以获取一小部分行(包括图像 blob)并将其直接写入本地文件系统的文件中。

from datasets import load_dataset

# Return as a Hugging Face dataset
ds = load_dataset(
    "lance-format/laion-1m",
    split="train",
    streaming=True
)

dir_name = "laion_samples"
Path(dir_name).mkdir(exist_ok=True)

for idx, row in enumerate(ds.take(3)):
    with open(f"{dir_name}/{idx}.jpg", "wb") as f:
        f.write(row["image"])

在这个示例中,image 列包含编码后的图像字节,因此您可以直接将其写入 .jpg 文件。

datasets API 目前不会将操作下推到 Lance 表,因此对于较大的数据集,速度可能会较慢。目前,直接使用 lance Python 包会获得更好的性能。请参阅 Hub 上的文档了解使用示例。

图像解码

默认情况下,当您迭代数据集时,图像会按顺序解码为 PIL.Images。不过,使用多线程解码可以显著提高数据集的处理速度

>>> 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 慢。

如果您不需要将图像解码为 PIL.Images,而是希望访问路径/字节,您可以禁用解码

>>> dataset = dataset.decode(False)

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

在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.