torchvision.datasets.ImageFolder
官方文档: ImageFolder
ImageFolder是一个通用的数据集加载API,继承自DataFolder,其要求数据集的排列如下所示
root/dog/xxx.png
root/dog/xxy.png
root/dog/[...]/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/[...]/asd932_.png
其包含一些属性
- root: 数据集的路径.
- transform: 对图片数据进行操作。输入为PIL图片并且返回一个transform之后的版本
- target_transform: 对类别标签数据进行操作。输入为类别并将其进行转换
注意: 虽然ImageFolder继承自DataFolder,但是其基类都来自torch.utils.data.Dataset,因此都可以用Dataset的一些方法,例如len(dataset)获取数据集的大小
ImageFolder用法
这里我们创建一个train
文件夹,里面数据集的格式如下
运行下面代码
import torchvision.datasets as dset
data_root = './train'
dataset = dset.ImageFolder(root="train")
print(dataset[0])
'''
(<PIL.Image.Image image mode=RGB size=300x280 at 0x7FE5F80B0590>, 0)
'''
可以看到如果我们没有加入transform信息,ImageFolder
整合后的数据类型为PIL,下面我们加上一个transform
trans = transforms.ToTensor()
dataset = dset.ImageFolder(root="train", transform=trans)
print(dataset[0])
'''
(tensor([[[0.1529, 0.1529, 0.1569, ..., 0.8118, 0.7922, 0.7882],
[0.1569, 0.1569, 0.1569, ..., 0.7961, 0.7804, 0.7725],
[0.1569, 0.1569, 0.1569, ..., 0.7804, 0.7765, 0.7765],
...,
[0.1216, 0.1137, 0.0980, ..., 0.0667, 0.0941, 0.1569],
[0.1216, 0.1137, 0.0980, ..., 0.0824, 0.0980, 0.1412],
[0.1216, 0.1098, 0.0941, ..., 0.1176, 0.0902, 0.0824]]]), 0)
'''
上述例子可以看到,PIL格式的图像在数据加载时已经变为tensor形式。同理如果需要对类别target进行处理,加入target_transform
即可
ImageFolder属性
除此之外,ImageFolder由于继承了DataFolder
类,所以包含他们的一些属性
- classes: 类别属性,根据类文件名来确定类别
- class_to_idx: 按顺序将这些类别索引为0,1,2…
- imgs: 返回所有文件夹中的文件路径即其类别
- targets: 返回文件夹中的文件类别
data_root = './train'
dataset = dset.ImageFolder(root="train")
print(dataset.classes) #根据分的文件夹的名字来确定的类别
print(dataset.class_to_idx) #按顺序为这些类别定义索引为0,1...
print(dataset.imgs) #返回从所有文件夹中得到的图片的路径以及其类别
print(dataset.loader)
print(dataset.extensions)
print(dataset.targets)
'''
['cat', 'dog']
{'cat': 0, 'dog': 1}
[('train/cat/cat.1.jpg', 0), ('train/cat/cat.2.jpg', 0), ('train/dog/dog.1.jpg', 1), ('train/dog/dog.2.jpg', 1)]
<function default_loader at 0x7fb2e911c320>
('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
[0, 0, 1, 1]
'''