基于PyTorch框架的图像分类之——1.加载图像数据的两种方式

系列文章目录

基于PyTorch框架的图像分类之——1.加载图像数据的两种方式

---

---

前言

PyTorch框架进行图像分类时候，我们这里以加载训练集为例来说明。加载训练集数据，实质是加载图像数据并赋予分类标签。

---

一、ImageFolder函数

用ImageFolder函数加载图像训练集的实质是，加载出图像数据，并对每个图片依据其在哪个文件夹中自动编码生成对应文件夹的标签。

其代码如下

import os
import torchvision
from torchvision import transforms,datasets


data_dir = "C:\\Users\\zijiancao\\Desktop\\where_am_i"

train_dir = 'train'

data_transform = transforms.Compose([
    transforms.ToTensor()
])
trainset = torchvision.datasets.ImageFolder(root=os.path.join(data_dir,train_dir), transform=data_transform                                  )
print(trainset.class_to_idx)

这里输出的是每个类别对应的编码

{‘CALsuburb’: 0, ‘PARoffice’: 1, ‘bedroom’: 2, ‘coast’: 3, ‘forest’: 4, ‘highway’: 5, ‘industrial’: 6, ‘insidecity’: 7, ‘kitchen’: 8, ‘livingroom’: 9, ‘mountain’: 10, ‘opencountry’: 11, ‘store’: 12, ‘street’: 13, ‘tallbuilding’: 14}

---

这里的root为存储训练集的地址，要求文件夹组织如下图：

---

训练集的文件夹中包含许多类别的文件

点开bedroom，里面都是照片

二、自定义加载的数据集

对于没有对类别编码要求，我们可以直接调用ImageFolder函数加载加载数据集，相应对于类别编码如第一个图所示。但是对于有类别对应编码要求，如第二个图所示。这时候需要我们自定义加载数据集了。