PyTorch ImageFolder自定义数据集

在PyTorch自定义数据集中，我们介绍了如何通过重写Dataset类来自定义数据集，但其实对于图像数据，自定义数据集有一个更简单的方法，那就是直接调用ImageFolder，它是torchvision.datasets里的函数。

ImageFolder介绍

ImageFolder假设所有的文件按文件夹保存，每个文件夹下存储同一个类别的图片，文件夹名为类名，其构造函数如下：

ImageFolder(root, transform=None, target_transform=None, loader=default_loader)

各参数含义：

root：在root指定的路径下寻找图片

transform：对PIL Image进行的转换操作，transform的输入是使用loader读取图片的返回对象

target_transform：对label的转换

loader：给定路径后如何读取图片，默认读取为RGB格式的PIL Image对象

label是按照文件夹名顺序排序后存成字典，即{类名:类序号(从0开始)}

示例

从kaggle官网下载dogsVScats的数据集（百度网盘的下载链接见文末），该数据集包含test1文件夹和train文件夹，train文件夹中包含12500张猫的图片和12500张狗的图片，图片的文件名中带序号：

cat.0.jpg
cat.1.jpg
cat.2.jpg
...
cat.12499.jpg
dog.0.jpg
dog.1.jpg
dog.2.jpg
...
dog.12499.jpg

假设我们希望把train文件夹中90%猫的图片和90%狗的图片作为训练集，剩下的10%作为验证集：

import os
import shutil
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset
from torchvision import transforms, datasets

# kaggle原始数据集在本地电脑的文件路径
original_dataset_dir = '/Users/wangpeng/Desktop/all/CS/Datasets/kaggle_dogs_cats/train'
total_num = int(len(os.listdir(original_dataset_dir)) / 2)
random_idx = np.array(range(total_num))
np.random.shuffle(random_idx)

# 待处理的数据集地址
base_dir = '/Users/wangpeng/Desktop/dogsVScats'
if not os.path.exists(base_dir):
    os.mkdir(base_dir)

# 训练集、验证集的划分
sub_dirs = ['train', 'validate']
animals = ['cats', 'dogs']
train_idx = random_idx[:int(total_num * 0.9)]
validate_idx = random_idx[int(total_num * 0.9):]
numbers = [train_idx, validate_idx]
for idx, sub_dir in enumerate(sub_dirs):
    dir = os.path.join(base_dir, sub_dir)
    if not os.path.exists(dir):
        os.mkdir(dir)
    for animal in animals:
        animal_dir = os.path.join(dir, animal)
        if not os.path.exists(animal_dir):
            os.mkdir(animal_dir)
        fnames = [animal[:-1] + '.{}.jpg'.format(i) for i in numbers[idx]]
        for fname in fnames:
            src = os.path.join(original_dataset_dir, fname)
            dst = os.path.join(animal_dir, fname)
            shutil.copyfile(src, dst)

        # 训练集、验证集的图片数目
        print(animal_dir + ' total images : %d' % (len(os.listdir(animal_dir))))

运行上面的程序，在我的电脑的桌面上将会有一个dogsVScats文件夹，其文件结构如下：

dogsVScats
    |
    |----train
    |         |     
    |         |---cats(包含11250张猫的图片)
    |         |---dogs(包含11250张狗的图片)
    |   
    |-----validate
              |
              |---cats(包含1250张猫的图片)
              |---dogs(包含1250张狗的图片)

接着我们就可以用ImageFolder创建数据集了，并把创建好的数据集放到DataLoader中：

data_transform = transforms.Compose([
    transforms.Resize(256),         # 把图片resize为256*256
    transforms.CenterCrop(224),     # 随机裁剪224*224
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])   # 标准化
])

train_dataset = datasets.ImageFolder(root='/Users/wangpeng/Desktop/dogsVScats/train', transform=data_transform)  # 标签为{'cats':0, 'dogs':1}
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True)

validate_dataset = datasets.ImageFolder(root='/Users/wangpeng/Desktop/dogsVScats/validate', transform=data_transform)  
validate_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)

我们可以测试一下，看一下train_loader可不可以用：

if __name__ == '__main__':
    image, label = iter(train_loader).next()  # iter()函数把train_loader变为迭代器，然后调用迭代器的next()方法
    sample = image[0].squeeze()
    sample = sample.permute((1, 2, 0)).numpy()
    sample *= [0.229, 0.224, 0.225]
    sample += [0.485, 0.456, 0.406]
    sample = np.clip(sample, 0, 1)
    plt.imshow(sample)
    plt.show()
    print('Label is: {}'.format(label[0].numpy()))

运行结果：

PyTorch ImageFolder自定义数据集