文章目录
什么是ImageFolder用代码来理解用法
什么是ImageFolder
torchvision.datasets.ImageFolder 是 PyTorch 中 torchvision 库提供的一个用于加载图像数据集的类，特别适用于处理按类别组织的图像数据集。通过使用 ImageFolder 类，你可以轻松地加载包含图像数据的文件夹，并且该类会自动根据文件夹结构解析出每个类别的图像数据
用代码来理解
假设我们有一个图像数据集，包含两个类别：“cat” 和 “dog”，每个类别各有一些图像。数据集的目录结构如下：
dataset/ ├── cat/ │ ├── cat001.jpg │ ├── cat002.jpg │ └── … └── dog/ ├── dog001.jpg ├── dog002.jpg └── …
在这个例子中，“dataset” 文件夹下有一个名为 “cat” 的子文件夹和一个名为 “dog” 的子文件夹，分别存放了 “cat” 类别和 “dog” 类别的图像。每个子文件夹中包含该类别的多张图像。
接下来，我们可以使用 ImageFolder 类来加载这个数据集，并查看返回的内容是什么：
import torch
from torchvision import datasets, transforms
# 定义数据转换
transform = transforms.Compose([
transforms.Resize((224, 224)), # 调整图像大小
transforms.ToTensor(), # 将图像转换为Tensor
])
# 加载数据集
dataset = datasets.ImageFolder('dataset', transform=transform)
# 查看数据集中的内容
print(dataset.classes) # 打印类别列表
print(dataset.class_to_idx) # 打印类别到索引的映射
print(dataset.imgs) # 打印图像路径列表
dataset.classes 返回的结果会是 [‘cat’, ‘dog’]，表示数据集中包含的类别 dataset.class_to_idx 返回的结果可能是 {‘cat’: 0, ‘dog’: 1}，表示每个类别对应的索引 而 dataset.imgs 则会返回一个列表，其中每个元素对应一个元组 (image_path, class_index)，包含了图像的路径和对应的类别索引。
用法
import torch
from torchvision import datasets, transforms
# 定义数据转换
transform = transforms.Compose([
transforms.Resize((224, 224)), # 调整图像大小
transforms.ToTensor(), # 将图像转换为Tensor
])
# 加载数据集
dataset = datasets.ImageFolder('path/to/dataset', transform=transform)
# 创建数据加载器
data_loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)
精彩内容
发表评论
2024-06-22 12:27:07回复