Comment utiliser la bibliothèque de traitement d'image basée sur PyTorch "Kornia"

introduction

J'utilise souvent PyTorch comme cadre d'apprentissage en profondeur, mais j'ai récemment découvert une bibliothèque de traitement d'image appelée «Kornia» basée sur PyTorch. J'ai donc étudié les fonctions de base et l'utilisation, donc je vais laisser cela comme un mémorandum.

Qu'est-ce que Kornia

Kornia est une bibliothèque de vision par ordinateur open source implémentée avec PyTorch comme backend. (Kornia GitHub)

It consists of a set of routines and differentiable modules to solve generic computer vision problems. At its core, the package uses PyTorch as its main backend both for efficiency and to take advantage of the reverse-mode auto-differentiation to define and compute the gradient of complex functions.

Il met en œuvre des opérations de traitement d'image de bas niveau similaires à OpenCV, telles que le filtrage, la conversion des couleurs et la conversion géométrique. Et comme PyTorch est le backend, il a l'avantage d'être plus sensible aux avantages du support GPU et de la différenciation automatique.

Installation et utilisation de base

L'installation peut être effectuée avec pip, etc. comme décrit dans README. (Dans ce cas, PyTorch sera entré automatiquement) pip install kornia

• Depuis le 20/06/2020, kornia 0.3.1 / pytorch 1.5.1 a été installé.

Vous aurez également besoin d'OpenCV, de matplotlib et de torchvision pour exécuter les différents tutoriels (https://kornia.readthedocs.io/en/latest/tutorials/index.html).

À titre d'exemple d'utilisation de base, le processus d'application du flou gaussien à une image est le suivant.

import kornia
import cv2

#Lecture d'image avec OpenCV
img_src = cv2.imread('./data/lena.jpg')
img_src = cv2.cvtColor(img_src, cv2.COLOR_BGR2RGB)

# torch.Convertir en tenseur
tensor_src = kornia.image_to_tensor(img_src, keepdim=False).float() # 1xCxHxW

# Gaussian Blur
gauss = kornia.filters.GaussianBlur2d((11, 11), (10.5, 10.5))
tensor_blur = gauss(tensor_src)

# OpenCV(numpy.ndarray)Revenir à l'image
img_blur = kornia.tensor_to_image(tensor_blur.byte())

# --> show [img_src | img_blur]

Autres exemples de traitement d'image

Un exemple de flou et de changement de couleur autre que celui ci-dessus est illustré ci-dessous.

# Box Blur
tensor_blur = kornia.box_blur(tensor_src, (9, 9))

# Median Blur
tensor_blur = kornia.median_blur(tensor_src, (5, 5))

# Adjust Brightness
tensor_brightness = kornia.adjust_brightness(tensor_src, 0.6)

# Adjust Contrast
tensor_contrast = kornia.adjust_contrast(tensor_src, 0.2)

# Adjust Gamma
tensor_gamma = kornia.adjust_gamma(tensor_src, gamma=3., gain=1.5)

# Adjust Saturation
tensor_saturated = kornia.adjust_saturation(tensor_src, 0.2)

# Adjust Hue
tensor_hue = kornia.adjust_hue(tensor_src, 0.5)

Combinaison avec torch.nn.Sequential

En résumant le traitement décrit ci-dessus avec nn.Sequential, vous pouvez écrire le prétraitement de l'image proprement. Ce qui suit est un exemple. De plus, ici, le traitement est effectué sur le GPU dans les prémisses d'un environnement où le GPU peut être utilisé.

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import kornia

class DummyDataset(Dataset):
    def __init__(self):
        self.data_index = range(100)

    def __len__(self):
        return len(self.data_index)

    def __getitem__(self, idx):
        # generate dummy image and label
        image = torch.rand(3, 240, 320)
        label = torch.randint(5, (1,))
        return image, label

device = torch.device('cuda')

dataset = DummyDataset()
loader = DataLoader(dataset, batch_size=16, shuffle=True)

transform = nn.Sequential(
    kornia.color.AdjustSaturation(0.2),
    kornia.color.AdjustBrightness(0.5),
    kornia.color.AdjustContrast(0.7),
)

for i, (images, labels) in enumerate(loader):
    print(f'iter: {i}, images: {images.shape}, labels: {labels.shape}')

    images = images.to(device) # -->Vers le GPU
    images_tr = transform(images) #Appliquer la transformation à l'image

    # training etc ...

Exemple d'utilisation de la différenciation automatique

Comme exemple d'utilisation de la différenciation automatique de PyTorch, je prendrai une partie du tutoriel total_variation_denoising.py (suppression totale du bruit de fluctuation).

# read the image with OpenCV
img: np.ndarray = cv2.imread('./data/doraemon.png')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
img = img + np.random.normal(loc=0.0, scale=0.1, size=img.shape)
img = np.clip(img, 0.0, 1.0)

# convert to torch tensor
noisy_image: torch.tensor = kornia.image_to_tensor(img).squeeze()  # CxHxW

# define the total variation denoising network
class TVDenoise(torch.nn.Module):
   def __init__(self, noisy_image):
       super(TVDenoise, self).__init__()
       self.l2_term = torch.nn.MSELoss(reduction='mean')
       self.regularization_term = kornia.losses.TotalVariation()
       # create the variable which will be optimized to produce the noise free image
       self.clean_image = torch.nn.Parameter(data=noisy_image.clone(), requires_grad=True)
       self.noisy_image = noisy_image

   def forward(self):
       return self.l2_term(self.clean_image, self.noisy_image) + 0.0001 * self.regularization_term(self.clean_image)

   def get_clean_image(self):
       return self.clean_image

tv_denoiser = TVDenoise(noisy_image)

# define the optimizer to optimize the 1 parameter of tv_denoiser
optimizer = torch.optim.SGD(tv_denoiser.parameters(), lr=0.1, momentum=0.9)

# run the optimization loop
num_iters = 500
for i in range(num_iters):
   optimizer.zero_grad()
   loss = tv_denoiser()
   if i % 25 == 0:
       print("Loss in iteration {} of {}: {:.3f}".format(i, num_iters, loss.item()))
   loss.backward()
   optimizer.step()

# convert back to numpy
img_clean: np.ndarray = kornia.tensor_to_image(tv_denoiser.get_clean_image())

Ici, noisy_image est passé à torch.nn.Parameter () pour faire de clean_image l'état initial. (Ceci est la cible de mise à jour par l'optimiseur) De plus, TotalVariation () de kornia est utilisé comme terme de régularisation.

• Le terme de régularisation de variation totale est un terme de pénalité lorsque la différence entre les pixels adjacents est grande, et il fonctionne pour supprimer les composants de bruit et lisser l'image.

Résumé

Nous avons étudié comment utiliser la bibliothèque de traitement d'images Kornia basée sur Pytorch, en nous concentrant sur des tutoriels. Il y aura probablement diverses fonctions utiles autres que celles que j'ai écrites cette fois. Il peut être utilisé non seulement dans la préparation d'image, mais également à l'avant du réseau neuronal, j'ai donc pensé que cela pourrait être utile lorsque vous souhaitez ajouter un traitement qui ne peut pas être géré par la torche / vision de la torche ordinaire seule.