Learn_PyTorch_2_数据加载与预处理

1 自定义数据集

本文利用的数据集是面部姿势数据集，内容为一个名为face_landmarks.csv和69张后缀为.jpg的面部图片。

其中，face_landmarkers.csv文件的内容格式如下：

文件名	x1	y1	···	x68	y68
0805personali01.jpg	27	83	···	84	134
1084239450_e76e00b7e7.jpg	70	236	···	128	312
···	···	···	···	···	···
personalpic.jpg	40	109	···	68	164

其中，第一列是每个样本的文件名，第2到第最后一列为每个样本图片的68个标记点的横纵坐标。

下面导入所有需要用到的包：

from __future__ import print_function, division # 为了兼容Python2.*的print和除法
import os # 文件路径需要用到
import torch # pytorch
import pandas as pd # 用来载入csv文件
from skimage import io, transform # 用来读取图片和对图片进行转换
import numpy as np	# numpy不解释
import matplotlib.pyplot as plt # 画图
from torch.utils.data import Dataset, DataLoader # 自定义数据集继承自它
from torchvision import transforms, utils # 一个高级计算机视觉库

plt.ion()   # 交互模式，为了让plt.show()不阻塞程序

1.1 从csv文件读取数据集名和标记

为了方便读取和转换，采用pandas进行数据载入：

landmarks_frame = pd.read_csv('data/faces/face_landmarks.csv') # 返回的是pandas的DataFrame数据结构

n = 10 # 第10+1张图片
img_name = landmarks_frame.iloc[n, 0] # .iloc根据下标索引，这里返回第10+1张图片的文件名
landmarks = landmarks_frame.iloc[n, 1:].as_matrix() # 将该图片的标记点坐标全部读出来
landmarks = landmarks.astype('float').reshape(-1, 2) # 将数据类型转换为float，并将shape转换为(68, 2)

print('Image name: {}'.format(img_name)) # 打印文件名
print('Landmarks shape: {}'.format(landmarks.shape)) # 打印保存标记点的数组的大小(68 * 2)
print('First 4 Landmarks: {}'.format(landmarks[:4])) # 打印前四个标记点

输出：

Image name:1878519279_f905d4f34e.jpg
Landmarks shape:(68, 2)
First 4 Landmarks:[[144. 178.]
 [145. 191.]
 [149. 205.]
 [157. 220.]]

1.2 显示样本图像和标记

def show_landmarks(image, landmarks):
    """显示图像和标记"""
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s = 10, marker = '.', c = 'r')
    plt.pause(0.001)
plt.figure()
show_landmarks(io.imread(os.path.join('data/faces/', img_name)), landmarks)
plt.show()

结果如下：

1.3 自定义数据集

torch.utils.data.Dataset是表示数据集的抽象类，自定义数据集时应继承Dataset并重写下面两个方法：

• __len__：使得我们可以通过len(dataset)获取数据集的大小
• __getitem__：能让我们以下标的方式索引数据集中的样本

下面创建一个数据集，通过__init__读取csv文件，加载图像的操作在调用__getitem__时进行以节约内存。

class FaceLandmarksDataset(Dataset):
    """数据集"""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        参数:
            csv_file (string): csv文件的路径
            root_dir (string): 数据集所在文件夹根路径
            transform (callable, optional): 变换
        """
        self.landmarks_frame = pd.read_csv(csv_file) # 读取csv文件
        self.root_dir = root_dir # 保存根路径
        self.transform = transform # 保存变换

    def __len__(self): # 返回数据集大小
        return len(self.landmarks_frame)

    def __getitem__(self, idx): # 返回单个样本
        img_name = os.path.join(self.root_dir,
                                self.landmarks_frame.iloc[idx, 0])
        image = io.imread(img_name)
        landmarks = self.landmarks_frame.iloc[idx, 1:].as_matrix()
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks} # 将样本组合为字典

        if self.transform: # 如果有变换则执行变换
            sample = self.transform(sample)

        return sample # 返回单个样本

下面尝试实例化FaceLandmarkDataset类：

face_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv', # 文件路径
                                    root_dir='data/faces/') # 根路径

fig = plt.figure()

for i in range(len(face_dataset)): # 迭代数据集
    sample = face_dataset[i] # 通过__getitem__获取单个样本

    print(i, sample['image'].shape, sample['landmarks'].shape) # 打印样本信息

    ax = plt.subplot(1, 4, i + 1) # 子图
    plt.tight_layout() # 紧密排列
    ax.set_title('Sample #{}'.format(i)) # 标题
    ax.axis('off') # 关闭坐标轴
    show_landmarks(**sample) # 将字典拆开传进去

    if i == 3: # 只迭代前四个样本
        plt.show()
        break

结果：

0 (324, 215, 3) (68, 2)
1 (500, 333, 3) (68, 2)
2 (250, 258, 3) (68, 2)
3 (434, 290, 3) (68, 2)

1.4 自定义变换

大多数神经网络都期望固定大小的图像，因此需要对图像进行缩放。此外，为了增加样本集大小，一个常用的数据增广操作是随机裁剪图像。下面是三个最常用的变换：

• Rescale：图像缩放，将所有样本大小调整为一致
• RandomCrop：随机裁剪，对样本集进行数据增广
• ToTensor：将numpy图像转换为torch图像

为了避免每次调试时都进行参数传递，因此将上述三个变换编写为可调用类而不是简单的函数。因此，需要实现三个类的__call__方法：

class Rescale(object):
    """将样本集缩放到指定大小

    参数:
        output_size (tuple or int): 输出大小。如果是tuple，则缩放到outuput_size大小；
        									 如果是int，则窄边缩放到output_size大小并保持宽长比不变
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple)) # 判断是否输入的是tuple/int
        self.output_size = output_size

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']

        h, w = image.shape[:2] # C * H * W 取后两维度为尺寸
        if isinstance(self.output_size, int): # int对应的缩放
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size # tuple对应的缩放

        new_h, new_w = int(new_h), int(new_w) # 取整

        img = transform.resize(image, (new_h, new_w)) # skimage中的transform.resize

        # numpys图像数据结构： h * w * c
        # torch图像数据结构:   c * h * w
        landmarks = landmarks * [new_w / w, new_h / h]

        return {'image': img, 'landmarks': landmarks} # 返回缩放结果


class RandomCrop(object):
    """随机缩放

    参数:
        output_size (tuple or int): 输出大小。如果是int，则输出大小为output_size * output_size
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int): # 如果是int，则输出大小为output_size * output_size
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, sample): # 调用
        image, landmarks = sample['image'], sample['landmarks']

        h, w = image.shape[:2] # numpy图像：h * w * c
        new_h, new_w = self.output_size

        top = np.random.randint(0, h - new_h) # 随机距离顶端
        left = np.random.randint(0, w - new_w) # 随机距离左边

        image = image[top: top + new_h, # 裁剪
                      left: left + new_w]

        landmarks = landmarks - [left, top] # 标记点坐标位移

        return {'image': image, 'landmarks': landmarks} # 返回裁剪结果


class ToTensor(object): 
    """将ndarray转换为tensor"""

    def __call__(self, sample): # 调用
        image, landmarks = sample['image'], sample['landmarks']

        # 交换轴序
        # numpy图像数据结构 : H x W x C
        # torch图像数据结构 : C X H X W
        image = image.transpose((2, 0, 1))
        return {'image': torch.from_numpy(image),
                'landmarks': torch.from_numpy(landmarks)} # 返回转换结果

torchvision.transforms.Compose是一个可调用类，他可以将上述三个变换进行合并操作，下面组合三个变换并应用到样本中：

scale = Rescale(256) 
crop = RandomCrop(128) 
composed = transforms.Compose([Rescale(256),
                               RandomCrop(224)])
# 上述两个变换首先将图像的窄边变换到256，然后从图像中随机裁剪出一个224*224的子图

# 将上述变换应用于样本
fig = plt.figure()
sample = face_dataset[65] # 第66个样本

# 对该样本分别执行缩放、裁剪和缩放并裁剪三种变换 
for i, tsfrm in enumerate([scale, crop, composed]):
    transformed_sample = tsfrm(sample)

    ax = plt.subplot(1, 3, i + 1)
    plt.tight_layout()
    ax.set_title(type(tsfrm).__name__)
    show_landmarks(**transformed_sample)

plt.show()

结果：

结果分别显示了第66个样本执行缩放、裁剪、缩放+裁剪的结果。

下面将我们自定义的变换作用到自定义的数据集上：

transformed_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv',
                                           root_dir='data/faces/',
                                           transform=transforms.Compose([
                                               Rescale(256),
                                               RandomCrop(224),
                                               ToTensor()
                                           ]))
# 实例化类的时候并没有对样本进行变换
# 当我们索引样本时调用__getitem__时才执行变换

for i in range(len(transformed_dataset)):
    sample = transformed_dataset[i] # 索引样本的时候才执行的变换

    print(i, sample['image'].size(), sample['landmarks'].size())

    if i == 3: # 仅迭代前四个样本
        break

结果：

0 torch.Size([3, 224, 224]) torch.Size([68, 2])
1 torch.Size([3, 224, 224]) torch.Size([68, 2])
2 torch.Size([3, 224, 224]) torch.Size([68, 2])
3 torch.Size([3, 224, 224]) torch.Size([68, 2])

1.5 数据集迭代器

上面已经基本完成了一个数据集的预处理，但仍缺失了一些更加高效的技巧，如批量化加载数据、随机打乱数据和多线程加载数据等。

torch.utils.data.DataLoader是一个提供了以上功能的迭代器：

# 数据集迭代器
dataloader = DataLoader(transformed_dataset, batch_size=4,
                        shuffle=True, num_workers=4)
# 如果出错令num_workers = 0

def show_landmarks_batch(sample_batched):
    """批量化显示样本"""
    images_batch, landmarks_batch = \
            sample_batched['image'], sample_batched['landmarks']
    batch_size = len(images_batch) # 批量大小
    im_size = images_batch.size(2) # 图像大小
    grid_border_size = 2 # 网格边界的宽度

    grid = utils.make_grid(images_batch) # 画个网格
    plt.imshow(grid.numpy().transpose((1, 2, 0))) 

    for i in range(batch_size):
        plt.scatter(landmarks_batch[i, :, 0].numpy() + i * im_size + (i + 1) * grid_border_size,
                    landmarks_batch[i, :, 1].numpy() + grid_border_size,
                    s=10, marker='.', c='r')

        plt.title('Batch from dataloader')

for i_batch, sample_batched in enumerate(dataloader): # 利用数据集迭代器进行迭代（每次四个样本）
    print(i_batch, sample_batched['image'].size(),
          sample_batched['landmarks'].size())

    # 迭代四次共16个样本
    if i_batch == 3:
        plt.figure()
        show_landmarks_batch(sample_batched)
        plt.axis('off')
        plt.ioff()
        plt.show()
        break

2 利用torchvision包构建数据集

torchvision是一个完全独立于pytorch的一个图像工具包，主要包含以下几个部分：

• torchvision.datasets：包含几个常用的图像数据集。
• torchvision.models：包含常用的模型如：AlexNet, VGG, ResNet等。
• torchvision.transforms：包含常用的图像操作，如缩放，裁剪等。这里主要用到这个部分。
• torchvision.utils：一些工具，如张量保存、生成网格以可视化批量图像等。

利用torchvision中的ImageFolder可以很方便的生成数据集，这里要求图像按如下方式组织：

root/train/class1/xxx.png
root/train/class1/xxy.png
root/train/class1/xxz.png
·
·
·
root/train/calss2/aaa.png
root/train/class2/aab.png
root/train/class2/aac.png
·
·

下面以hymenoptera_data数据集为例，其目录结构如下：

hymenoptera_data/train/ants/aaa.png
                            ....png
hymenoptera_data/train/bees/bbb.png
                            ....png

hymenoptera_data/val/ants/ccc.png
                          ....png
hymenoptera_data/val/bees/ddd.png
                          ....png  

构建数据集如下：

# 训练集和验证集的变换略有不同
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224), # 随机大小、宽长比裁剪图片，然后resize到244*244
        transforms.RandomHorizontalFlip(), # 以默认概率0.5进行水平翻转
        transforms.ToTensor(), # 转换为张量
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        # 归一化，对应mean和std是imageNet计算出来的
    ]),
    'val': transforms.Compose([
        transforms.Resize(256), # resize
        transforms.CenterCrop(224), # 中心裁剪
        transforms.ToTensor(), #  转换为张量
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = 'data/hymenoptera_data' # 数据集路径
# 下面得到的数据集是一个字典：
# image_datasets = {'train':trainsets, 'val':valsets}
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
# 下面得到的数据集迭代器也是一个字典：
# dataloaders = {'train':trainLoader, 'val':valLoader}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
# 数据集大小也是字典：dataset_sizes = {'train':trainSize, 'val':valSize}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

上述代码中关于训练集和验证集的变换有所不同。对于训练集，这里采用的随机裁剪和随机水平翻转的方式进行了数据增广。（数据增广不是直接增加了训练集的大小，而是在每次加载样本时对样本进行上述随机变换从而间接增加的训练集的规模）