ImageFolder 是 PyTorch 中 torchvision.datasets 模块提供的一个常用类，用于从文件夹中加载图像数据。它是一种非常方便的方式来加载按文件夹结构组织的图像数据集。这个类能够自动将文件夹中的子目录作为标签，并且将其中的图像文件加载为 PyTorch 张量。
1. 基本概念
ImageFolder 假定数据集的文件夹结构是这样的：
root/
├── class_1/
│ ├── img1.jpg
│ ├── img2.jpg
│ └── ...
├── class_2/
│ ├── img1.jpg
│ ├── img2.jpg
│ └── ...
├── class_3/
│ ├── img1.jpg
│ ├── img2.jpg
│ └── ...
└── ...
每个子文件夹（例如 class_1、class_2）代表一个类别，文件夹中的图像文件属于该类别。ImageFolder 会根据每个文件夹的名称来为图像分配标签（例如，class_1 对应标签 0，class_2 对应标签 1，依此类推）。
2. ImageFolder 的使用
创建 ImageFolder 对象
你可以通过指定数据集所在的根目录来创建 ImageFolder 对象。例如：
from torchvision import datasets, transforms
# 数据集的根目录
root = 'path/to/your/dataset'
# 数据预处理的转换操作
transform = transforms.Compose([
transforms.Resize((128, 128)), # 将图像调整为 128x128 大小
transforms.ToTensor(), # 将图像转换为 Tensor
])
# 创建 ImageFolder 数据集对象
dataset = datasets.ImageFolder(root=root, transform=transform)
ImageFolder 类的关键参数
- root: 数据集的根目录，通常是包含所有类别文件夹的上级目录。
- transform: 用于数据增强和预处理的
transform操作。它会被应用到每张图像上。例如，你可以使用transforms.Resize()、transforms.ToTensor()等。 - target_transform: 用于标签的变换操作，类似于
transform，但作用于标签（类别）。 - loader: 默认情况下，
ImageFolder使用PIL图像加载器加载图像。你可以传入自定义的加载函数。
ImageFolder 返回的数据结构
ImageFolder 类返回一个包含两部分的元组：
- 图像: 图像数据通常是一个
PIL图像对象或者经过transform转换后的 PyTorch 张量。 - 标签: 图像的标签，通常是一个整数，表示图像所属的类别。标签是根据文件夹名称生成的，
class_1的标签为 0，class_2的标签为 1，依此类推。
3. 如何使用 ImageFolder
访问图像和标签
通过索引，你可以获取 ImageFolder 中的图像和标签：
image, label = dataset[0]
image是经过预处理后的 PyTorch 张量（例如，(C, H, W)的张量）。label是图像对应的类别标签（整数）。
使用 DataLoader 迭代数据
为了方便批量加载数据，你通常会将 ImageFolder 与 DataLoader 结合使用：
from torch.utils.data import DataLoader
# 创建 DataLoader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 迭代 DataLoader 获取数据
for images, labels in dataloader:
print(images.shape) # 输出形状，例如 (32, 3, 128, 128)
print(labels) # 输出对应的标签
4. 示例代码
假设我们有以下文件夹结构：
data/
├── dogs/
│ ├── dog1.jpg
│ ├── dog2.jpg
│ └── ...
├── cats/
│ ├── cat1.jpg
│ ├── cat2.jpg
│ └── ...
我们可以使用 ImageFolder 来加载这个数据集，并进行处理：
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 设置图像预处理操作
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
])
# 创建 ImageFolder 数据集对象
dataset = datasets.ImageFolder(root='data', transform=transform)
# 创建 DataLoader 对象
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 迭代数据
for images, labels in dataloader:
print(images.shape) # 例如 (32, 3, 128, 128)
print(labels) # 例如 tensor([0, 1, 0, 1, ..., 0, 1])，0 表示狗，1 表示猫
在这个例子中，ImageFolder 会根据文件夹 dogs 和 cats 的名称自动分配标签。对于 dogs 文件夹中的图像，标签是 0；对于 cats 文件夹中的图像，标签是 1。
5. 总结
ImageFolder是一个非常方便的类，可以自动从文件夹结构中加载图像，并为每个类别生成标签。- 它适用于经典的图像分类任务，其中图像按类别存储在不同的文件夹中。
- 你可以通过
transform参数自定义图像预处理流程（例如调整大小、转换为张量等），并通过DataLoader实现批量加载和数据迭代。