пятница, 18 января 2019 г.

PyTorch: загрузка и обработка данных изображений

Много усилий в решении любой проблемы машинного обучения идет на подготовку данных. PyTorch предоставляет множество инструментов для упрощения загрузки данных. В этом посте мы рассмотрим, как загружать и обрабатывать/дополнять данные из нетривиального набора данных.

Чтобы выполнять примеры из этого поста, убедитесь, что следующие пакеты установлены:

  • scikit-image - для ввода-вывода изображений и преобразований
  • pandas - для более простого разбора csv

# -*- coding: utf-8 -*-
from __future__ import print_function, division
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils

# игнорируем предупреждения
import warnings
warnings.filterwarnings("ignore")

plt.ion()   # интерактивный режим

Набор данных, с которым мы собираемся иметь дело, это набор лицевых поз. Это означает, что лицо помечается так:

В целом, 68 различных ориентиров отмечены для каждого лица.

Примечание. Загрузите набор данных отсюда и сделайте так, чтобы они располагались в каталоге названном data/faces/.

Набор данных поставляется с CSV-файлом с аннотациями, который выглядит следующим образом:

image_name,part_0_x,part_0_y,part_1_x,part_1_y,part_2_x, ... ,
part_67_x,part_67_y
0805personali01.jpg,27,83,27,98, ... 84,134
1084239450_e76e00b7e7.jpg,70,236,71,257, ... ,128,312

Давайте быстро прочитаем 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:].as_matrix()
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]))

Вывод:

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

Давайте напишем простую вспомогательную функцию, чтобы показать изображение и его ориентиры, и использовать его для показа образца.

# -*- coding: utf-8 -*-
def show_landmarks(image, landmarks):
    """Показывает изображение с ориентирами"""
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', 
                c='r')
    plt.pause(0.001)  # pause a bit so that plots are updated

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


Класс Dataset

torch.utils.data.Dataset - абстрактный класс, представляющий набор данных. Ваш пользовательский набор данных должен наследовать Dataset и переопределять следующие методы:

  • __len__, чтобы len(dataset) возвращал размер набора данных.
  • __getitem__ для поддержки индексации, так что dataset[i] может использоваться для получения i-го экземпляра

Давайте создадим dataset класс для нашего набора данных лицевых ориентиров. Мы будем читать csv в __init__, но оставим чтение изображений для __getitem__. Это эффективно для памяти, потому что все изображения не сохраняются в памяти сразу, но читаются по мере необходимости.

Образец нашего набора данных будет словарем (dict) {'image': image, 'landmarks': landmarks}. Наш набор данных примет необязательный аргумент transform, так что любая необходимая обработка может быть наносима на образец. Мы увидим пользу transform в следующем разделе.

# -*- coding: utf-8 -*-

class FaceLandmarksDataset(Dataset):
    """Face Landmarks dataset."""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Args:
            csv_file (string): Путь к csv файлу с аннотациями.
            root_dir (string): Каталог со всеми изображениями.
            transform (callable, optional): Необязательный transform 
                который будет применен к экземпляру.
        """
        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:].as_matrix()
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

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

        return sample

Давайте создадим экземпляр этого класса и пройдемся по образцам данных. Мы распечатаем размеры первых 4 образцов и покажем их ориентиры.

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)
    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)


Преобразования

Одна из проблем, которые мы видим из вышеизложенного, заключается в том, что экземпляры не одинакового размера. Большинство нейронных сетей ожидают изображения фиксированного размера. Поэтому нам нужно написать некоторый код предварительной обработки. Давайте создадим три преобразования:

  • Rescale: масштабировать изображение
  • RandomCrop: обрезать изображение случайным образом. Это увеличение данных.
  • ToTensor: преобразовать отдельные изображения в torch изображения (нам нужно поменять оси).

Мы напишем их как вызываемые классы вместо простых функций, так что параметры преобразования не нужно передавать каждый раз, когда класс будет вызываться. Для этого нам просто нужно реализовать метод __call__ и если требуется, метод __init__. Затем мы можем использовать преобразование следующим образом:

tsfm = Transform(params)
transformed_sample = tsfm(sample)

Обратите внимание, как эти преобразования должны применяться как к изображению, так и к ориентирам.

# -*- coding: utf-8 -*-

class Rescale(object):
    """Масштабирует изображение экземпляра в заданный размер.

    Args:
        output_size (tuple or int): Требуемый размер вывода. 
              Если tuple, вывод соответствует 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 и w поменялись местами для ориентиров, 
        # потому что для изображений,
        # x и y оси - это оси 1 и 0 соответственно
        landmarks = landmarks * [new_w / w, new_h / h]

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


class RandomCrop(object):
    """Обрезает случайным образом изображение экземпляра.

    Args:
        output_size (tuple or 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_w]

        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 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)}


Совмещение преобразований

Теперь мы применим преобразования к образцу.

Допустим, мы хотим изменить масштаб более короткой стороны изображения до 256 и затем случайным образом обрезать квадрат размером 224 от него. Мы хотим совместить Rescale и RandomCrop преобразования. torchvision.transforms.Compose - это простой вызываемый класс, который позволяет нам сделать это.

scale = Rescale(256)
crop = RandomCrop(128)
composed = transforms.Compose([Rescale(256),
                               RandomCrop(224)])

# Применяем каждую из трансформаций выше к экземпляру.
fig = plt.figure()
sample = face_dataset[65]
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()


Итерация по набору данных

Давайте соединим все это, чтобы создать набор данных с трансформациями.
Подводя итог, каждый раз, когда выбирается экземпляр этого набора данных:

  • Изображение читается из файла на лету
  • Преобразования применяются к считанному изображению
  • Поскольку одно из преобразований является случайным, данные увеличиваются при выборке

Мы можем выполнить итерацию по созданному набору данных с помощью
for i in range цикла, как и раньше.

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

Вывод:

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])

Однако мы теряем множество функций, используя простой цикл for для перебора данных. В частности, нам не хватает:

  • Пакетирования данных
  • Перемешивания данных
  • Загрузки данных параллельно, используя multiprocessing воркеры (многопроцессорные обработчики).

torch.utils.data.DataLoader - это итератор, который предоставляет все эти функции. Параметры, используемые ниже, должны быть ясны. Один параметр, представляющий интерес - это collate_fn. Вы можете указать, как именно образцы должны быть упакованным с использованием collate_fn. Тем не менее, сортировка по умолчанию должна работать хорошо для большинства случаев использования.

# -*- coding: utf-8 -*-

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


# Helper function to show a batch
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 = 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,
                    landmarks_batch[i, :, 1].numpy(),
                    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())

    #  просматриваем 4й пакет и останавливаемся
    if i_batch == 3:
        plt.figure()
        show_landmarks_batch(sample_batched)
        plt.axis('off')
        plt.ioff()
        plt.show()
        break

Вывод:

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])


Резюме: torchvision

В этом посте мы расмотрели, как писать и использовать наборы данных, преобразования и загрузчики данных. Пакет torchvision предоставляет несколько общеиспользуемых наборов данных и трансформаций. Возможно, вам даже не придется писать собственные классы. Один из наиболее общеиспользуемых наборов данных, доступных в 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, такие как RandomHorizontalFlip, Scale. Вы можете использовать их для написания загрузчика следующим образом:

import torch
from torchvision import transforms, datasets

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])
    ])
hymenoptera_dataset = datasets.ImageFolder(
                                    root='hymenoptera_data/train',
                                    transform=data_transform)
dataset_loader = torch.utils.data.DataLoader(hymenoptera_dataset,
                                    batch_size=4, shuffle=True,
                                    num_workers=4)


Читайте также: