基于Python的人工智能应用案例系列（13）：自定义数据集、DataLoader和Transforms

        在许多人工智能应用中，处理不同形式和格式的数据是必不可少的。本文将通过一个实际案例，展示如何在PyTorch中编写自定义数据集（Dataset）、DataLoader和数据转换（Transforms）。我们将使用面部标记（Facial Landmarks）数据集，并演示如何从图像中提取面部特征点。

1. 加载与可视化数据

        首先，我们需要从CSV文件中读取数据，并可视化图像和标记点。面部标记数据集包含面部图像以及每张脸上的68个关键点坐标。

import os
import pandas as pd
from skimage import io
import matplotlib.pyplot as plt
import numpy as np

# 读取面部标记数据
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.reshape(-1, 2)

# 显示图像和标记点
def show_landmarks(image, landmarks):
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')

plt.figure()
image = io.imread(os.path.join('../data/faces/', img_name))
show_landmarks(image, landmarks)
plt.show()

2. 自定义Dataset类

        在PyTorch中，Dataset是一个抽象类，用户可以继承该类并重写其中的方法来自定义数据集。我们将创建一个自定义的FaceLandmarksDataset类，以便从图像和CSV文件中读取数据。

from torch.utils.data import Dataset

class FaceLandmarksDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        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:].values.reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

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

        return sample

3. 定义Transforms

        通常情况下，神经网络期望输入的图像具有固定大小。因此，我们需要在数据进入模型之前对其进行一些预处理，包括图像缩放和裁剪。为此，我们可以定义几个自定义的Transforms，如Rescale、RandomCrop和ToTensor。

from skimage import transform

class Rescale(object):
    def __init__(self, output_size):
        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, self.output_size
        img = transform.resize(image, (new_h, new_w))
        landmarks = landmarks * [new_w / w, new_h / h]
        return {'image': img, 'landmarks': landmarks}

class RandomCrop(object):
    def __init__(self, output_size):
        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, 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):
    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']
        image = image.transpose((2, 0, 1))
        return {'image': torch.from_numpy(image), 'landmarks': torch.from_numpy(landmarks)}

4. 应用Transforms

        我们可以组合这些Transforms，并将它们应用到数据集中。

from torchvision import transforms

# 组合Transforms
composed = transforms.Compose([Rescale(256), RandomCrop(224), ToTensor()])

# 应用到数据集
transformed_dataset = FaceLandmarksDataset(csv_file='../data/faces/face_landmarks.csv',
                                           root_dir='../data/faces/',
                                           transform=composed)

# 显示处理后的图像和标记
for i in range(4):
    sample = transformed_dataset[i]
    print(i, sample['image'].shape, sample['landmarks'].shape)

5. 使用DataLoader进行批处理

    DataLoader可以帮助我们更高效地加载数据，支持批处理、数据打乱和多线程加载。

from torch.utils.data import DataLoader

dataloader = DataLoader(transformed_dataset, batch_size=4, shuffle=True, num_workers=0)

# 显示批处理中的图像和标记点
def show_landmarks_batch(sample_batched):
    images_batch, landmarks_batch = sample_batched['image'], sample_batched['landmarks']
    plt.imshow(utils.make_grid(images_batch).numpy().transpose((1, 2, 0)))
    for i in range(len(images_batch)):
        plt.scatter(landmarks_batch[i, :, 0].numpy() + i * images_batch.size(2), landmarks_batch[i, :, 1].numpy(), s=10, c='r')

for i_batch, sample_batched in enumerate(dataloader):
    show_landmarks_batch(sample_batched)
    if i_batch == 3:
        break

结语

        在本文中，我们通过一个面部标记数据集的案例，学习了如何在PyTorch中编写自定义的Dataset类、应用Transforms进行数据预处理，并利用DataLoader实现数据批处理。我们首先介绍了如何加载并可视化图像和标记点，随后定义了几个关键的自定义Transforms，如Rescale、RandomCrop和ToTensor，用于图像缩放、裁剪和张量转换。最后，通过DataLoader我们实现了数据的批量加载和处理，从而为神经网络的训练准备好数据。

        这种灵活的数据处理方式为我们处理各种非结构化数据提供了强大的工具，尤其适合在实际项目中面对多样化的数据类型。接下来，我们可以将这些处理好的数据集输入到深度学习模型中，进一步进行训练、优化和推理，开启人工智能的下一步探索。

如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！

欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。

谢谢大家的支持！