В мире обработки изображений часто встречаются изображения с различными соотношениями сторон. Однако при работе со сверточными нейронными сетями (CNN) или другими моделями, требующими входных данных квадратной формы, нам часто необходимо преобразовать эти изображения в идеальные квадраты. В этой статье блога мы рассмотрим различные методы приведения изображений к квадратным размерам с помощью PyTorch. Я расскажу вам о каждом методе, используя разговорный язык, и попутно приведу примеры кода. Итак, приступим!

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

import torch
def pad_to_square(image):
    height, width, _ = image.shape
    max_dim = max(height, width)
    pad_height = max_dim - height
    pad_width = max_dim - width
    padded_image = torch.zeros((max_dim, max_dim, 3))
    padded_image[:height, :width, :] = image
    return padded_image

Метод 2: заполнение средним цветом
Вместо заполнения нулями мы также можем дополнить изображение средним значением цвета исходного изображения. Это помогает сохранить визуальное содержание изображения, делая его квадратным. Вот пример реализации:

import torch
def pad_to_square_mean(image):
    height, width, _ = image.shape
    max_dim = max(height, width)
    pad_height = max_dim - height
    pad_width = max_dim - width
    mean_color = torch.mean(image, dim=(0, 1))
    padded_image = torch.zeros((max_dim, max_dim, 3))
    padded_image[:height, :width, :] = image
    padded_image[height:, width:, :] = mean_color
    return padded_image

Метод 3: обрезка по центру
Другой подход к преобразованию изображений в квадраты — обрезка по центру. Мы посчитаем разницу между шириной и высотой, а затем обрежем изображение одинаково с обеих сторон. Вот пример реализации:

import torch
def pad_to_square_center_crop(image):
    height, width, _ = image.shape
    max_dim = max(height, width)
    pad_height = max_dim - height
    pad_width = max_dim - width
    top = pad_height // 2
    bottom = pad_height - top
    left = pad_width // 2
    right = pad_width - left
    padded_image = torch.nn.functional.pad(image, (left, right, top, bottom))
    return padded_image

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

import torch
from torchvision import transforms
def pad_to_square_rescale(image):
    height, width, _ = image.shape
    max_dim = max(height, width)
    rescale_transform = transforms.Resize((max_dim, max_dim))
    rescaled_image = rescale_transform(image)
    return rescaled_image

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