I often use PyTorch as a framework for deep learning, but I recently learned about an image processing library called "Kornia" that is based on PyTorch. So I investigated the basic functions and usage, so I will leave it as a memorandum.
Kornia is an open source computer vision library implemented with PyTorch as the 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.
It implements low-level image processing operations similar to OpenCV, such as filtering, color conversion, and geometric conversion. And because PyTorch is the back end, it has the advantage of being more susceptible to the benefits of GPU support and automatic differentiation.
Installation can be done with pip etc. as described in the README. (In this case, PyTorch will be entered automatically)
pip install kornia
You will also need OpenCV, matplotlib, and torchvision to run the tutorials (https://kornia.readthedocs.io/en/latest/tutorials/index.html).
As an example of basic usage, the process of applying Gaussian Blur to an image is as follows.
import kornia
import cv2
#Image reading with OpenCV
img_src = cv2.imread('./data/lena.jpg')
img_src = cv2.cvtColor(img_src, cv2.COLOR_BGR2RGB)
# torch.Convert to Tensor
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)Return to image
img_blur = kornia.tensor_to_image(tensor_blur.byte())
# --> show [img_src | img_blur]
In this way, the target processing is performed for torch.Tensor.
(By the way, kornia.filters.GaussianBlur2d inherits torch.nn.Module)
An example of blurring processing and color change other than the above is shown below.
# 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)
By summarizing the processing described above with nn.Sequential, you can write the preprocessing of the image neatly. The following is an example. Also, here, processing is performed on the GPU on the premise of an environment where the GPU can be used.
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) # -->To GPU
images_tr = transform(images) #Apply transform to image
# training etc ...
As an example of using PyTorch's automatic differentiation, I will quote a part from the tutorial total_variation_denoising.py (total variation noise removal).
# 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())
Here, noisy_image is passed to torch.nn.Parameter () to make clean_image the initial state. (This is the update target by optimizer) In addition, kornia's TotalVariation () is used as a regularization term.
We investigated how to use the Pytorch-based image processing library Kornia, focusing on tutorials. There are likely to be various useful functions other than those I wrote this time. It can be used not only in image preparation but also in the forward of neural networks, so I thought it might be useful when you want to add processing that cannot be handled by ordinary torch / torchvision alone.
Recommended Posts