torchvision.datasets.ImageFolder

torchvision , , 评论

torchvision.datasets.ImageFolder

官方文档: ImageFolder

ImageFolder是一个通用的数据集加载API，继承自DataFolder，其要求数据集的排列如下所示

1
2
3
4
5
6
7
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文件夹，里面数据集的格式如下

运行下面代码

1
2
3
4
5
6
7
8
9
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

1
2
3
4
5
6
7
8
9
10
11
12
13
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: 返回文件夹中的文件类别
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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]
'''

harryteaustc

computer vision