pytorch中torchvision模块下ImageFolder的简单理解与实际运用

pytorch中torchvision模块下ImageFolder的简单理解与实际运用

ImageFolder函数定义

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

root: 图片总目录，子层级为各类型对应的文件目录。
transform: 对PIL image进行转换操作，transform输入的loader读取图片返回的对象
target_transform: 对label进行变换
loader: 指定加载图片的函数，默认操作是读取为RGB格式的PIL image对象 

使用ImageFolder需要的目录结构

基本结构如下：

具体参考：

说明：
根目录root : my_data
子分类目录: classA , classB
classA与classB目录之下才是具体图片

注意点：

1. root路径必须指向到分类文件目录同级。也就是my_data下，而不能指向具体分类目录。即不能指向my_data/classA
2. 通过ImageFolder读取后，所有类别的所有图片都会被一同读取进来，而不是分类读取。
   len(data)的长度为7，即所有图片。
3. ImageFolder会将所有分类给一个类别标签索引值，该类别标签索引值从0开始
   data.class_to_idx 结果为：{‘classA’: 0, ‘classB’: 1}
4. ImageFolder生成的结果是一个列表，而该列表里的每个元素是一个元组，每个元组对应一张图片信息。具体内容为：(PIL图片,类别标签索引)

读取实例

from torchvision import transforms,utils
from torchvision import datasets
import torch
import matplotlib.pyplot as plt

#对图片进行随机裁剪256的大小
trans = transforms.Compose([
    transforms.RandomResizedCrop(256)
])

#使用ImageFolder读取图片
data = datasets.ImageFolder('my_data',transform=trans)

#读取了所有类别下的所有图片
len(data)
#7

#获取类别
data.classes
#['classA', 'classB']

#类别与标签索引
data.class_to_idx
#{'classA': 0, 'classB': 1}

#图片对应类别标签
data.targets
[0, 0, 0, 0, 1, 1, 1]

#每个元素是元组
data[0]
#(<PIL.Image.Image image mode=RGB size=256x256 at 0x23836184828>, 0)


#查看所有图片路径与类别标签
data.imgs
#
[('my_data\\classA\\a.png', 0),
 ('my_data\\classA\\b.png', 0),
 ('my_data\\classA\\c.png', 0),
 ('my_data\\classA\\d.png', 0),
 ('my_data\\classB\\m1.png', 1),
 ('my_data\\classB\\m2.png', 1),
 ('my_data\\classB\\m3.png', 1)]

查看单张图片

plt.imshow(data[0][0])