pytorch Dataset 的ImageFolder

转自来源：CSDN
原文：https://blog.csdn.net/TH_NUM/article/details/80877435

介绍常用的pytorch Dataset 里面的ImageFolder，实现和https://blog.csdn.net/TH_NUM/article/details/80877196很相似。
ImageFolder假设所有的文件按文件夹保存好，每个文件夹下面存贮同一类别的图片，文件夹的名字为分类的名字。

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

root : 在指定的root路径下面寻找图片
transform: 对PIL Image进行转换操作,transform 输入是loader读取图片返回的对象
target_transform :对label进行变换
loader: 指定加载图片的函数，默认操作是读取PIL image对象

例子：

from torchvision.datasets import ImageFolder

dataset=ImageFolder('data/dogcat_2/')

#对应文件夹的label
print(dataset.class_to_idx)

输出：

{'cat': 0, 'dog': 1}

#所有图片的路径和对应的label
print(dataset.imgs)

输出：

[(‘data/dogcat_2/cat/cat.12484.jpg’, 0), (‘data/dogcat_2/cat/cat.12485.jpg’, 0), (‘data/dogcat_2/cat/cat.12486.jpg’, 0), (‘data/dogcat_2/cat/cat.12487.jpg’, 0), (‘data/dogcat_2/dog/dog.12496.jpg’, 1), (‘data/dogcat_2/dog/dog.12497.jpg’, 1), (‘data/dogcat_2/dog/dog.12498.jpg’, 1), (‘data/dogcat_2/dog/dog.12499.jpg’, 1)]

#没有任何转变，所有返回的还是PIL Image对象
print(dataset[0][1]) #第二维度为1 ，表示label
print(dataset[0][0]) #第二维度为0，表示图片数据

输出:

< PIL.Image.Image image mode=RGB size=497x500 at 0x7F25F3D31E10>

完整代码：

from torchvision.datasets import ImageFolder
import torch
from torchvision import transforms

#加上transforms
normalize=transforms.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
transform=transforms.Compose([
    transforms.RandomReSizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(), #将图片转换为Tensor,归一化至[0,1]
    normalize
])

dataset=ImageFolder('data/dogcat_2/',transform=transform)

#输出第0张图片的大小
print(dataset[0][0].size())

输出：

torch.Size([3, 224, 224])