Pytorch基础(三)数据集加载及预处理

目录

下载数据集及显示样本

数据集类

建立数据集类及显示部分样本

数据变换

后记


python提供了许多工具简化数据加载,使代码更具可读性。经常用到的包有scikit-image、pandas等,本文通过相关包进行数据加载和预处理相关简要介绍。

此处(提取码:ilqy)下载数据集,数据存于"data/faces/"的目录中。这个数据集实际上是imagenet数据集标注为face的图片当中在dlib面部检测(dlib's pose estimation)表现良好的图片。下面以该数据集为例,对数据加载即预处理进行简要介绍。

下载数据集及显示样本

下面为下载数据集及显示其中某一样本的相关代码:

#!/usr/bin/env torch
# -*- coding:utf-8 -*-
# @Time  : 2021/2/3, 23:54
# @Author: Lee
# @File  : test.py

import os, torch
import pandas as pd
from skimage import io, transform  # 用于图像的IO和变换
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils, datasets

import warnings
warnings.filterwarnings("ignore")

plt.ion()  # interactive mode

# # 读取数据集  将csv中的标注点数据读入(N,2)数组中,其中N是特征点的数量
landmarks_frame = pd.read_csv("data/faces/face_landmarks.csv")

n = 65
img_name = landmarks_frame.iloc[n, 0]
landmarks = landmarks_frame.iloc[n, 1:].values
landmarks = landmarks.astype('float').reshape(-1, 2)


print('Image name: {}'.format(img_name))
print('Landmarks shape: {}'.format(landmarks.shape))
print('First 4 Landmarks: {}'.format(landmarks[:4]))


# 展示一张图片和它对应的标注点
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()
plt.pause(0)

打印结果如下:

Image name: person-7.jpg
Landmarks shape: (68, 2)
First 4 Landmarks: [[32. 65.]
 [33. 76.]
 [34. 86.]
 [34. 97.]]

显示图形如下:

数据集类

torch.utils.data.Dataset是表示数据集的抽象类,因此自定义数据集应继承Dataset并覆盖以下方法*__len__实现len(dataset)返回数据集的尺寸。*__getitem__用来获取一些索引数据,例如dataset[i]中的(i)。

建立数据集类及显示部分样本

为面部数据集创建一个数据集类。在__init__中读取csv的文件内容,在__getitem__中读取图片。这么做是为了节省内存空间。只有在需要用到图片的时候才读取它而不是一开始就把图片全部放到内存中。

数据样本将按这样一个字典{'image':image, 'landmarks':landmarks}组织。该数据类将添加一个可选参数transform以方便对样本进行预处理,代码如下:

#!/usr/bin/env torch
# -*- coding:utf-8 -*-
# @Time  : 2021/2/4, 0:13
# @Author: Lee
# @File  : dataset_class.py

import os, torch
import pandas as pd
from skimage import io, transform  # 用于图像的IO和变换
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils, datasets

import warnings
warnings.filterwarnings("ignore")

plt.ion()  # interactive mode


# 展示一张图片和它对应的标注点
def show_landmarks(image, landmarks):
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')
    plt.pause(0.001)


# 数据集类
# 数据样本按这样一个字典{'image': image, 'landmarks': landmarks}组织。
# 添加一个可选参数transform 以方便对样本进行预处理
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)
        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:]
        landmarks = np.array([landmarks])
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample


# 获取图片并可视化部分图片
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]
    print(i, sample['image'].shape, sample['landmarks'].shape)
    ax = plt.subplot(1, 4, i+1)
    ax.set_title('Sample #{}'.format(i))
    ax.axis('off')
    show_landmarks(**sample)

    if i == 3:
        plt.show()
        break
plt.pause(0)

运行结果如下:

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

数据变换

通过上面的例子可知数据集中的图片并不是同样的尺寸。绝大多数神经网络都假定图片的尺寸相同。因此需要做预处理。这里以三个转换*Rescale(缩放图片),*RandomCrop(对图片进行随机剪裁),*ToTensor(把numpy格式的图片转换为torch格式图片,需要交换坐标轴)。

可以将它们写成可调用的类的形式而不是简单的函数,这样就不需要每次调用时传递一遍参数。只需要实现__call__方法,必要的时候实现__init__方法。

把这些整合起来以创建一个带组合转换的数据集。每次这个数据集被采样时,*即使地从文件中读取图片*对读取的图片应用转换*,由于其中一步是随机的(Randmpcrop),数据有所增强,现在可以用循环来对所有创建的数据执行同样的操作。

但是,对所有数据简单地使用for循环牺牲了很多功能,尤其是*批量处理数据(指定batch_size)*打乱数据(shuffle置True)*使用多线程(multiprocessingworker)并加载数据。

torch.utils.data.DataLoader是一个提供了上述所有这些功能的迭代器。下面使用的参数必须是清楚的。一个值得关注的参数是collate_fn,可以通过它来决定如何对数据进行批量处理,但是绝大多数情况下默认值就能运行良好。

代码如下:

#!/usr/bin/env torch
# -*- coding:utf-8 -*-
# @Time  : 2021/2/4, 0:28
# @Author: Lee
# @File  : data_preprogress.py

import os, torch
import pandas as pd
from skimage import io, transform  # 用于图像的IO和变换
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils, datasets

import warnings
warnings.filterwarnings("ignore")

plt.ion()  # interactive mode


# 展示一张图片和它对应的标注点
def show_landmarks(image, landmarks):
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')
    plt.pause(0.001)


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)
        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:]
        landmarks = np.array([landmarks])
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample


class Rescale(object):
    """
    将样本中的图像重新缩放到给定大小
    Args:
    output_size(tuple或int):所需的输出大小。如果是元组,则输出为
    与output_size匹配。如果是int,则匹配较小的边缘到output_size保持横纵比相同
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size =output_size

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']
        h, w = image.shape[:2]
        if isinstance(self.output_size, 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

        new_h, new_w = int(new_h), int(new_w)
        img = transform.resize(image, (new_h, new_w))
        # h and w are swapped for landmarks because for images,
        # x and y axes are axis 1 and 9 respectively
        landmarks = landmarks * [new_w / w, new_h / h]
        return {'image': img, 'landmarks': landmarks}


class RandomCrop(object):
    """
    随机裁剪样本中的图像
    Args:
        output_size(tuple或int):所需的输出大小,如果是int, 方形裁剪是
    """
    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            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]
        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_h]
        landmarks = landmarks - [left, top]

        return {'image': image, 'landmarks': landmarks}


class ToTensor(object):
    """将样本中的ndarrays转换为Tensors"""
    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']
        """
        交换颜色轴原因
        numpy包的图片时H*W*C 而torch包的图片是 C*H*W
        """
        image = image.transpose((2, 0, 1))
        return {'image': torch.from_numpy(image),
                'landmarks': torch.from_numpy(landmarks)}


# 辅助功能:显示批次
def show_landmark_batch(sample_batched):
    """show image with landmarks for a batch of samples"""
    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')


if __name__ == '__main__':
    # 获取图片并可视化部分图片
    # 数据变换及torchvision.transforms.Compose组合操作
    face_dataset = FaceLandmarksDataset(csv_file='data/faces/face_landmarks.csv',
                                        root_dir='data/faces/')
    scale = Rescale(256)
    crop = RandomCrop(128)
    composed = transforms.Compose([Rescale(256),
                                   RandomCrop(224)])
    # 在样本上应用上述的变换
    fig = plt.figure()
    sample = face_dataset[65]
    print('数据变换及torchvision.transforms.Compose组合操作')
    for i, tsfrm in enumerate([scale, crop, composed]):
        transformed_sample = tsfrm(sample)
        ax = plt.subplot(1, 3, i + 1)
        plt.tight_layout()
        ax.axis('off')
        ax.set_title(type(tsfrm).__name__)
        show_landmarks(**transformed_sample)
    plt.show()
    plt.pause(0.5)

    # 迭代数据集
    transformed_dataset = FaceLandmarksDataset(csv_file="data/faces/face_landmarks.csv",
                                               root_dir="data/faces/",
                                               transform=transforms.Compose([Rescale(256), RandomCrop(224), ToTensor()]))

    for i in range(len(transformed_dataset)):
        sample = transformed_dataset[i]
        print(i, sample['image'].size(), sample['landmarks'].size())

        if i == 3:
            break

    dataloader = DataLoader(transformed_dataset, batch_size=4,
                            shuffle=True, num_workers=4)
    print('迭代数据集,batch_size=4')
    for i_batch, sample_batched in enumerate(dataloader):
        print(i_batch, sample_batched['image'].size(),
              sample_batched['landmarks'].size())

        if i_batch == 3:
            plt.figure()
            show_landmark_batch(sample_batched)
            plt.axis('off')
            plt.ioff()
            plt.show()
            break
    plt.pause(0)





运行结果如下:

数据变换及torchvision.transforms.Compose组合操作
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])
迭代数据集,batch_size=4
0 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
1 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
2 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
3 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])

后记

上面的例子用函数实现了数据的部分预处理操作,主要包括使用数据集类(datasets),转换(transform)和数据加载器(DataLoader)。torchvision包提供了畅通的数据集类datasets和转换transforms,可能并不需要我们自己构造这些类。torchvision中海油一个更常用的数据集类ImageFolder。它假定了数据集是以如下方式构造的:

root/ants/xxx.png
root/ants/xxy.jpeg
root/ants/xxz.png
...
root/bees/123.jpg
root/bees/nsdf3.png
root/bees/asd932_.png

其中'ants','bees'等是分类标签。在PIL.Image中可以使用类似的转换transform,例如RandHorizontalFlip,Scale。利用这些可以按如下方式创建一个数据集加载器(dataloader),以hymenoptera_data(提取码:2rvf)数据集为例:

#!/usr/bin/env torch
# -*- coding:utf-8 -*-
# @Time  : 2021/2/4, 1:06
# @Author: Lee
# @File  : data_transform.py

import os, torch
import pandas as pd
from skimage import io, transform  # 用于图像的IO和变换
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils, datasets

import warnings
warnings.filterwarnings("ignore")

plt.ion()  # interactive mode

data_transform = transforms.Compose([
    transforms.RandomSizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0,225])
    ])
hymenpotera_dataset = datasets.ImageFolder(root='data/hymenoptera_data/train',
                                           transform=data_transform)
dataset_loader = torch.utils.data.DataLoader(hymenpotera_dataset,
                                              batch_size=4, shuffle=True,
                                              num_workers=4)

print(dataset_loader)

打印结果如下:

<torch.utils.data.dataloader.DataLoader object at 0x000001C37596F608>

Debug可获得如下窗口:

 

  • 3
    点赞
  • 8
    收藏
    觉得还不错? 一键收藏
  • 4
    评论
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值