Python, Программирование

Kornia — Python библиотека для обработки изображений в задачах CV.

Время прочтения: 4 мин.

В этой статье я хотел бы познакомить читателей с библиотекой для ЯП python — Kornia, имеющей богатый функционал в области computer vision. Библиотека написана с использованием pytorch, в ее основе лежат готовые решения, такие как torchvision, PIL, skimage, tf.image, OpenCV. В Kornia реализована возможность выполнения вычислений не только с использованием CPU, но и GPU, что делает фреймворк довольно производительным.

            Библиотека содержит в своем составе ряд компонентов, пробежимся по основным:

  1. kornia.enhance – предоставляет набор функций, для улучшения изображений:
  • adjust_brightness – управляет яркостью изображения.
  • adjust_contrast – управление контрастностью.
  • adjust_saturation – насыщенность изображения.
  • adjust_gamma – гамма-контроль.
  • adjust_hue – управление оттенками.

Ниже приведен код, и результат работы метода adjust_contrast:

import matplotlib.pyplot as plt
import cv2
import kornia as K
import numpy as np
import torch
import torchvision
from PIL import Image


# Функция для считывания изображения с диска и преобразования его в torch-массив для последующей обработки:
def load_torch_image(fname):
    img = K.image_to_tensor(cv2.imread(fname), False).float() / 255
    img = K.color.bgr_to_rgb(img)
    return img

# Функция отрисовки и вывода на экран torch-массива при помощи MatPlotLib:
def imshow(tensor):
    image = torchvision.utils.make_grid(tensor, nrow=2)
    image = K.utils.tensor_to_image(image)
    plt.imshow(image)
    plt.axis('off')
    plt.show()

# Функция для сохранения torch-массива как изображения на диск:
def save_tensor_as_image(path, tensor):
    image = torchvision.utils.make_grid(tensor, nrow=2)
    image = K.utils.tensor_to_image(image)
    im = Image.fromarray(np.uint8(image*255))
    im.save(path)



fname0 = 'kot_sneg.png'

img1 = load_torch_image(fname0)

img2 = K.enhance.adjust_contrast(img1, 1.15)

save_tensor_as_image('kot_sneg_0.png', img2)

2. kornia.color – функции для работы со цветом:

  • bgr_to_rgb – приводит формат цветопредставления BGR к RGB.
  • rgb_to_grayscale – RGB к черно-белому.
  • rgb_to_hsv – привести RGB изображение к цветовой модели HSV.

Результат работы kornia.color.rgb_to_grayscale():

img1 = load_torch_image(fname0)
img2 = K.color.rgb_to_grayscale(img1)
save_tensor_as_image('kot_sneg_0.png', img2)

3. kornia.filters – набор методов для фильтрации изображений, а так же функции определения границ на изображениях:

  • box_blur
  • blur_pool2d
  • gaussian_blur2d
  • max_blur_pool2d
  • median_blur
  • motion_blur
    (Различные виды размытия, прямоугольное, по Гауссу, движением и т.д.)
  • spatial_gradient
  • sobel
  • laplacian
  • canny
    (Функции определения границ на изображениях)
    kornia.filters.canny() в деле:
img1 = load_torch_image(fname0)
img2 = K.filters.canny(img1)[0]
save_tensor_as_image('kot_sneg_0.png', 1. - img2.clamp(0., 1.))
  1. kornia.morphology – набор инструментов для математической морфологии:
  • morphology.dilation
  • morphology.erosion
  • morphology.opening
  • morphology.closing
  • morphology.gradient
  • morphology.bottom_hat
  • morphology.top_hat

(Это функции эрозии, наращивания, размыкания и т.д.) kornia.morphology.dilation():

img1 = load_torch_image(fname0)
kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [1, 0, 1]]).to('cpu') #Структурирующий элемент
img2 = K.morphology.dilation(img1, kernel)
save_tensor_as_image('kot_sneg_0.png', img2)

2. kornia.geometry – методы для работы с геометрией, в том числе с искажениями перспективы:

  • get_perspective_transform – функция вычисления трансформации перспективы.
  • warp_perspective – применение трансформации к изображению.
  • rotate – поворот изображения.
  • get_rotation_matrix2d – расчет матрицы множественного поворота.
  • warp_affine – применить вращение из матрицы.
  • rescale – изменение размера (есть antialiasing).

Ниже код и результат связки из get_perspective_transform() и warp_perspective():

img1 = load_torch_image(fname0)

points_src = torch.tensor([[
    [0., 0.], [1079., 0.], [1079., 959.], [0., 959.]
]])

points_dst = torch.tensor([[
    [30., 160.], [1040., 15.], [1040., 490.], [30., 930.]
]])

M = K.geometry.get_perspective_transform(points_src, points_dst)

img_warp = K.geometry.warp_perspective(img1.float(), M, dsize=(960, 1080))

save_tensor_as_image('kot_sneg_0.png', img_warp)

3. kornia.augmentation – инструменты аугментации:

  • ColorJitter
  • RandomAffine
  • RandomPerspective

Ниже код и результат его работы:

img1 = load_torch_image(fname0)

aug = K.augmentation.AugmentationSequential(
    K.augmentation.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0),
    K.augmentation.RandomAffine(360, [0.1, 0.1], [0.7, 1.2], [30., 50.], p=1.0),
    K.augmentation.RandomPerspective(0.5, p=1.0),
    data_keys=['input'],
    return_transform=False,
    same_on_batch=False,
)
for i in range(4):
    img2 = aug(img1)
    save_tensor_as_image(f'kot_sneg_{i}.png', img2)

Вышеперечисленными методами функционал библиотеки не ограничивается. Специалисты, занятые разработкой ПО в области computer vision могут найти в ней множество инструментов, которые помогут сэкономить время на написание кода. Например, в kornia реализован алгоритм Connected-component labeling (CCL) для сегментации мелких объектов на изображении, на базе этого фреймворка легко реализовать обучение модели Torch с автоматической аугментацией датасета и разделением данных на батчи, при этом есть возможность выбора между CPU или GPU.

Kornia может стать удобной надстройкой над классическими инструментами работы с изображениями, она активно поддерживается разработчиками и имеет емкую, удобную документацию.(https://kornia.readthedocs.io/en/latest/)

Советуем почитать