torchvision.datasets.ImageFolder

1. 函数说明

ImageFolder是一个通用的数据加载器，它要求我们以下面这种格式来组织数据集的训练、验证或者测试图片。

2. 文件夹存放格式

|--data
	|--train
		|--cat
		|--dog
	|--valid
		|--cat
		|--dog

3. ImageFolder参数

dataset = ImageFolder(root=path, transform=train_transforms,
					  target_transform=None,
					  loader=<function default_loader>,
					  is_valid_file=None)

• root : train 或者valid文件夹的地址
• transform: 输入图片转换
• target_transform: 输出图片转换
• is_vaild_file :接受Image文件的路径并检查该文件是否为验证集文件夹的函数

4. 代码

# -*- coding: utf-8 -*-
# @Project: zc
# @Author: zc
# @File name: ImageFolder_test
# @Create time: 2022/1/10 8:48

# 1.导入相关数据库
import torchvision.transforms
from torchvision.datasets import ImageFolder
from torchvision import transforms
import os

# 2.定义图片转换方式
train_transforms = torchvision.transforms.Compose([
	transforms.RandomResizedCrop(400),
	transforms.ToTensor()
])

# 3. 定义地址
path = os.path.join(os.getcwd(), 'data', 'train')

# 4. 将文件夹数据导入
dataset = ImageFolder(root=path, transform=train_transforms)

# 5.打印数据集描述
print(f'dataset={dataset}')

# 6.dataset里面存储的是[image_to_tensor,classes],即[图片张量形式，类别号]
# 7.dataset[0][0],指的是第一个训练样本中图片的张量大小
print(f'dataset[0][0]={dataset[0][0].shape}')

# 8.定义的是第一个样本图片的类别，cat->0;dog->1
print(f'dataset[0][1]={dataset[0][1]}')

# 9.打印数据集中所有的类别，返回类别列表
print(f'dataset.classes={dataset.classes}')

# 10.将类别列表进行编号，比如 cat=0,dog=1,,,依次类推
print(f'dataset.class_to_idx={dataset.class_to_idx}')

# 11.打印所有的图片和图片对应的类别
print(f'dataset.imgs={dataset.imgs}')

5. 结果

dataset=Dataset ImageFolder
    Number of datapoints: 5
    Root location: D:\zc\data\train
    StandardTransform
Transform: Compose(
               RandomResizedCrop(size=(400, 400), scale=(0.08, 1.0), ratio=(0.75, 1.3333), interpolation=bilinear)
               ToTensor()
           )
dataset[0][0]=torch.Size([3, 400, 400])
dataset[0][1]=0
dataset.classes=['cat', 'dog']
dataset.class_to_idx={'cat': 0, 'dog': 1}
dataset.imgs=[('D:\\zc\\data\\train\\cat\\cat1.jpg', 0), ('D:\\zc\\data\\train\\cat\\cat2.jpg', 0), ('D:\\zc\\data\\train\\cat\\cat3.jpg', 0), ('D:\\zc\\data\\train\\dog\\dog1.jpg', 1), ('D:\\zc\\data\\train\\dog\\dog2.jpg', 1)]