torchvision.ImageFolder
ImageFolder
一个通用的数据加载器,可以对我们自己的数据集进行处理,类似于DataSet加载CIFAR-10,COCO这样已经封装好的数据集。
ImageFolder
继承自 DatasetFolder
,因此可以重写相同的方法来自定义数据集,并且在DatasetFolder
中对所有的图片都进行了Image.convert("RGB")
处理,因此可以同时处理三通道、单通道图像
类原型:
CLASS torchvision.datasets.ImageFolder(root: str,
transform: ~typing.Optional[~typing.Callable] = None,
target_transform: ~typing.Optional[~typing.Callable] = None,
loader: ~typing.Callable[[str], ~typing.Any] = <function default_loader>,
is_valid_file: ~typing.Optional[~typing.Callable[[str], bool]] = None)
构造函数:
ImageFolder(root, transform=None, target_transform=None, loader=default_loader)
参数:
root (string) :根目录路径,标签的上一层目录,下图标签为cat与dog,上一层目录则到./train
transform (callable, optional) : 对PIL Image进行的转换操作,transform的输入是使用loader读取图片的返回对象
target_transform (callable, optional) : 对label(target)的转换
loader (callable, optional) : 给定路径后如何读取图片,默认读取为RGB格式的PIL Image对象
is_valid_file : 一个获取Image文件路径并检查文件是否是有效文件的函数(用于检查损坏的文件)
可以很方便的对如下结构的数据集进行处理:
其中cat、dog分别是训练数据集以及测试数据集的label(target)
成员变量:
-
classes
根据分的文件夹的名字来确定的类别target,如[‘cat’, ‘dog’]
-
class_to_idx
按顺序为这些类别定义索引为0,1…,如{‘cat’: 0, ‘dog’: 1}。
扫描二维码关注公众号,回复: 16514028 查看本文章 -
imgs
返回从所有文件夹中得到的图片的路径以及其类别,一个列表,列表中的每个元素都是一个(img-path, class_index)的元组
如:
[('./data_1/train/cat/0.jpg', 0), ('./data_1/train/cat/1.jpg', 0), ('./data_1/train/dog/0.jpg', 1), ('./data_1/train/dog/1.jpg', 1)]
实例程序:
from torchvision.datasets import ImageFolder
from torchvision import transforms
ROOT_TRAIN = r'./data_1/train'
train_transform = transforms.Compose([
transforms.ToTensor(),
])
train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
print(train_dataset.classes) # 得到所有的target
print(train_dataset.class_to_idx)
print(train_dataset.targets)
print(train_dataset[0])
print(train_dataset.imgs)
运行结果:
['cat', 'dog']
{
'cat': 0, 'dog': 1}
[0, 0, 1, 1]
(tensor([[[0.7961, 0.7961, 0.8000, ..., 0.9412, 0.9373, 0.9333],
[0.7961, 0.7961, 0.8000, ..., 0.9412, 0.9373, 0.9373],
[0.7961, 0.7961, 0.8000, ..., 0.9451, 0.9412, 0.9373],
...,
[0.6000, 0.6000, 0.6000, ..., 0.0078, 0.0078, 0.0078],
[0.5961, 0.5961, 0.5961, ..., 0.0039, 0.0039, 0.0039],
[0.5922, 0.5922, 0.5961, ..., 0.0039, 0.0039, 0.0039]],
[[0.6431, 0.6431, 0.6471, ..., 0.7882, 0.7843, 0.7804],
[0.6431, 0.6431, 0.6471, ..., 0.7882, 0.7843, 0.7843],
[0.6431, 0.6431, 0.6471, ..., 0.7922, 0.7882, 0.7843],
...,
[0.4784, 0.4784, 0.4784, ..., 0.0078, 0.0078, 0.0078],
[0.4745, 0.4745, 0.4745, ..., 0.0039, 0.0039, 0.0039],
[0.4706, 0.4706, 0.4745, ..., 0.0039, 0.0039, 0.0039]],
[[0.3412, 0.3412, 0.3451, ..., 0.4784, 0.4745, 0.4706],
[0.3412, 0.3412, 0.3451, ..., 0.4784, 0.4745, 0.4745],
[0.3412, 0.3412, 0.3451, ..., 0.4824, 0.4784, 0.4745],
...,
[0.2157, 0.2157, 0.2157, ..., 0.0000, 0.0000, 0.0000],
[0.2118, 0.2118, 0.2118, ..., 0.0000, 0.0000, 0.0000],
[0.2078, 0.2078, 0.2118, ..., 0.0000, 0.0000, 0.0000]]]), 0)
[('./data_1/train/cat/0.jpg', 0), ('./data_1/train/cat/1.jpg', 0), ('./data_1/train/dog/0.jpg', 1), ('./data_1/train/dog/1.jpg', 1)]