Park Lam的每日分享

Python实用工具img2dataset：大规模图像数据集高效构建指南

作者：

parklam

在

一、img2dataset库核心概述

img2dataset是一款专为大规模图像数据集构建设计的Python工具库，其核心用途是从图像URL列表中批量下载、处理并存储图像数据，广泛应用于计算机视觉领域的模型训练数据准备工作。该库的工作原理是通过多线程/多进程并行处理URL队列，支持断点续传、图像格式转换、分辨率调整等功能，同时能够生成配套的元数据文件，便于后续的数据管理与模型训练。

在优缺点方面，img2dataset的优势十分突出：一是并行处理机制大幅提升下载效率，能够轻松应对百万级以上的URL列表；二是支持多种输出格式（如webdataset、files、parquet等），适配不同的训练框架需求；三是内置图像过滤功能，可自动剔除损坏、低分辨率的无效图像。其缺点主要在于对网络环境要求较高，大规模下载时容易受带宽限制，且部分高级功能需要依赖额外的第三方库。该库采用Apache License 2.0开源协议，允许商用与二次开发，完全满足开发者的使用需求。

二、img2dataset安装方法

img2dataset支持通过Python包管理工具pip直接安装，同时也可以从GitHub源码编译安装，两种方式分别适用于不同的使用场景，以下是详细的安装步骤。

2.1 pip快速安装

这是最简便的安装方式，适用于大多数用户，只需在命令行中执行以下命令即可完成安装：

pip install img2dataset

安装完成后，可通过以下Python代码验证是否安装成功：

import img2dataset
print(f"img2dataset版本：{img2dataset.__version__}")

运行上述代码，如果控制台输出对应的版本号，说明安装成功；若出现ModuleNotFoundError，则需要检查pip环境是否配置正确，或尝试升级pip后重新安装。

2.2 源码编译安装

如果需要使用最新的开发版本，或者对源码进行自定义修改，可以选择从GitHub克隆源码并编译安装，步骤如下：

  1. 克隆GitHub仓库
git clone https://github.com/rom1504/img2dataset.git
cd img2dataset
  1. 安装依赖并编译
pip install -r requirements.txt
pip install -e .

这种安装方式的优势在于可以随时通过git pull获取最新的功能更新，适合对功能有定制化需求的开发者。

三、img2dataset基础使用教程

img2dataset的使用方式分为命令行调用和Python脚本调用两种，其中脚本调用的灵活性更高，便于嵌入到自动化数据处理流程中。本节将以Python脚本调用为主，结合实例讲解核心功能的使用方法。

3.1 核心参数说明

在使用img2dataset之前，需要先了解其核心参数的含义，这些参数决定了数据下载与处理的行为，关键参数如下表所示：

| 参数名称 | 数据类型 | 作用说明 | 默认值 |
|-|-|-|–|
| url_list | str | 存储图像URL的文件路径或文本内容 | 无（必填） |
| output_format | str | 输出格式，可选webdataset/files/parquet等 | webdataset |
| output_folder | str | 输出文件的存储目录 | dataset |
| thread_count | int | 并行下载的线程数 | 256 |
| image_size | int | 图像缩放后的目标分辨率 | 256 |
| resize_only_if_bigger | bool | 是否仅当原图大于目标分辨率时才缩放 | True |
| skip_reencode | bool | 是否跳过图像重新编码 | True |
| save_additional_columns | list | 需要保存的额外元数据列 | [] |

3.2 从URL列表下载图像（基础实例）

本实例将演示如何从一个包含图像URL的文本文件中批量下载图像，并保存为webdataset格式。

步骤1：准备URL列表文件

首先创建一个名为urls.txt的文本文件，每行存储一个图像URL和对应的元数据（如标签），格式如下：

https://example.com/image1.jpg label1
https://example.com/image2.jpg label2
https://example.com/image3.jpg label3

其中，URL与元数据之间用空格分隔，元数据可以根据需求添加多列。

步骤2：编写Python下载脚本

创建名为download_images.py的Python文件，代码如下：

from img2dataset import download

# 配置下载参数
params = {
    "url_list": "urls.txt",  # URL列表文件路径
    "output_folder": "my_image_dataset",  # 输出目录
    "output_format": "webdataset",  # 输出格式
    "thread_count": 128,  # 并行线程数，根据机器性能调整
    "image_size": 512,  # 图像缩放至512x512
    "resize_only_if_bigger": True,  # 仅缩放大于512的图像
    "skip_reencode": False,  # 重新编码为JPEG格式
    "save_additional_columns": ["label"],  # 保存标签列作为元数据
    "number_sample_per_shard": 1000,  # 每个分片存储1000张图像
    "retries": 3,  # 下载失败时重试次数
}

# 执行下载任务
download(**params)

步骤3：运行脚本并查看结果

在命令行中执行以下命令运行脚本：

python download_images.py

脚本运行后，会在当前目录下生成my_image_dataset文件夹，结构如下：

my_image_dataset/
├── 00000.tar
├── 00001.tar
└── ...

每个.tar文件是一个数据分片，包含1000张图像及其元数据，可直接用于PyTorch、TensorFlow等框架的模型训练。

3.3 直接使用URL列表字符串（进阶实例）

除了从文件读取URL列表，还可以直接将URL列表以字符串的形式传入参数，适用于动态生成URL的场景，代码示例如下：

from img2dataset import download

# 动态生成URL列表字符串
url_str = """https://example.com/img1.jpg cat
https://example.com/img2.jpg dog
https://example.com/img3.jpg bird
"""

# 配置参数
params = {
    "url_list": url_str,  # 直接传入URL字符串
    "output_folder": "dynamic_dataset",
    "output_format": "files",  # 以单个文件形式存储
    "image_size": 256,
    "thread_count": 64,
}

# 执行下载
download(**params)

该脚本运行后，dynamic_dataset文件夹下会按类别生成子文件夹，并存储对应的图像文件，适合需要人工查看图像的场景。

3.4 图像过滤与质量控制

img2dataset内置了图像质量过滤功能，可以自动剔除无效图像，例如损坏的文件、分辨率过低的图像等。以下是添加过滤条件的脚本示例：

from img2dataset import download

params = {
    "url_list": "urls.txt",
    "output_folder": "filtered_dataset",
    "output_format": "parquet",
    "image_size": 384,
    "min_image_size": 128,  # 剔除宽度或高度小于128的图像
    "max_image_area": 1000000,  # 剔除面积超过100万像素的图像
    "timeout": 10,  # 下载超时时间（秒）
    "verify_hash": False,  # 是否验证图像哈希值
    "skip_downloaded": True,  # 跳过已下载的图像（断点续传）
}

download(**params)

通过设置min_image_size和max_image_area参数，可以精准控制保留的图像质量，避免低质量数据影响模型训练效果。

四、img2dataset高级应用案例

本节将结合实际应用场景，讲解img2dataset的高级用法，包括与其他数据处理库的结合、大规模分布式下载等。

4.1 与Pandas结合处理元数据

在实际项目中，图像的元数据通常存储在CSV文件中，我们可以使用Pandas读取CSV文件，提取URL和元数据，再传递给img2dataset进行下载。以下是完整的案例代码：

import pandas as pd
from img2dataset import download

# 1. 使用Pandas读取CSV元数据文件
df = pd.read_csv("metadata.csv")
# 假设CSV包含列：url, label, category
print(f"元数据文件共包含 {len(df)} 条记录")

# 2. 将DataFrame转换为img2dataset支持的URL字符串格式
url_list = []
for idx, row in df.iterrows():
    url = row["url"]
    label = row["label"]
    category = row["category"]
    # 格式：URL 标签 类别
    url_list.append(f"{url} {label} {category}")
url_str = "\n".join(url_list)

# 3. 配置下载参数
params = {
    "url_list": url_str,
    "output_folder": "pandas_dataset",
    "output_format": "webdataset",
    "thread_count": 256,
    "image_size": 512,
    "save_additional_columns": ["label", "category"],  # 保存多列元数据
}

# 4. 执行下载
download(**params)

该案例适用于元数据较为复杂的场景，通过Pandas可以灵活地筛选、清洗元数据，再传递给img2dataset进行批量下载。

4.2 分布式大规模数据集下载

当需要处理千万级以上的URL列表时，单台机器的性能可能无法满足需求，此时可以使用img2dataset的分布式下载功能，借助多台机器并行处理任务。核心思路是将URL列表分割为多个分片，分配给不同的机器分别下载，最后合并结果。

步骤1：分割URL列表

使用以下Python代码将大型URL文件分割为多个小文件：

def split_url_file(input_file, chunk_size=100000):
    """
    将URL文件分割为多个分片
    :param input_file: 输入URL文件路径
    :param chunk_size: 每个分片的记录数
    """
    with open(input_file, "r", encoding="utf-8") as f:
        lines = f.readlines()

    total_chunks = (len(lines) + chunk_size - 1) // chunk_size
    for i in range(total_chunks):
        start = i * chunk_size
        end = min((i+1)*chunk_size, len(lines))
        chunk_lines = lines[start:end]
        with open(f"urls_chunk_{i}.txt", "w", encoding="utf-8") as f_out:
            f_out.writelines(chunk_lines)
    print(f"分割完成，共生成 {total_chunks} 个分片")

# 分割URL文件，每个分片10万条记录
split_url_file("large_urls.txt", chunk_size=100000)

步骤2：多机器并行下载

将分割后的URL分片文件分别发送到不同的机器，每台机器运行以下下载脚本：

from img2dataset import download

# 替换为对应的分片文件名
chunk_file = "urls_chunk_0.txt"

params = {
    "url_list": chunk_file,
    "output_folder": f"dataset_chunk_0",
    "output_format": "webdataset",
    "thread_count": 256,
    "image_size": 512,
    "distributor": "multiprocessing",  # 使用多进程分发任务
}

download(**params)

步骤3：合并下载结果

所有机器下载完成后，将生成的数据集分片复制到同一目录下，即可得到完整的大规模图像数据集。

五、img2dataset常见问题与解决方案

在使用img2dataset的过程中，可能会遇到各种问题，以下是一些常见问题及其解决方案：

5.1 下载速度慢

  • 原因：线程数设置过低，或网络带宽不足。
  • 解决方案：适当增加thread_count参数的值（根据机器CPU核心数调整，建议设置为CPU核心数的4-8倍）；使用高速网络，或配置代理服务器。

5.2 大量图像下载失败

  • 原因：URL无效、目标服务器拒绝访问，或下载超时。
  • 解决方案：增加retries参数的值，提高重试次数；设置合理的timeout参数；下载完成后查看日志文件，剔除无效URL。

5.3 内存占用过高

  • 原因：并行线程数过多，导致内存资源耗尽。
  • 解决方案：降低thread_count参数的值；使用distributor="multiprocessing"参数，采用多进程替代多线程，减少内存占用。

六、相关资源链接

  • Pypi地址：https://pypi.org/project/img2dataset
  • Github地址：https://github.com/rom1504/img2dataset
  • 官方文档地址：https://github.com/rom1504/img2dataset/blob/main/docs/README.md

关注我，每天分享一个实用的Python自动化工具。

更多文章