Open in Colab

Rotate image using warp affine transform#

In this tutorial we are going to learn how to rotate an image using the kornia.gemetry components.

!pip install kornia
import cv2
from matplotlib import pyplot as plt
import numpy as np

import torch
import torchvision
import kornia as K
/home/docs/checkouts/ TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See
  from .autonotebook import tqdm as notebook_tqdm

We use OpenCV to load an image to memory represented in a numpy.ndarray

img_bgr: np.ndarray = cv2.imread('bennett_aden.png', cv2.IMREAD_COLOR)

Convert the numpy array to torch

x_img: torch.Tensor = K.utils.image_to_tensor(img_bgr)  # CxHxW
x_img = x_img[None,...].float() / 255.
x_img = K.color.bgr_to_rgb(x_img)
def imshow(input: torch.Tensor, size: tuple = None):
    out = torchvision.utils.make_grid(input, nrow=4, padding=5)
    out_np: np.ndarray = K.utils.tensor_to_image(out)

Define the rotation matrix#

# create transformation (rotation)
alpha: float = 45.0  # in degrees
angle: torch.tensor = torch.ones(1) * alpha

# define the rotation center
center: torch.tensor = torch.ones(1, 2)
center[..., 0] = x_img.shape[3] / 2  # x
center[..., 1] = x_img.shape[2] / 2  # y

# define the scale factor
scale: torch.tensor = torch.ones(1, 2)

# compute the transformation matrix
M: torch.tensor = K.geometry.get_rotation_matrix2d(center, angle, scale)  # 1x2x3

Apply the transformation to the original image#

_, _, h, w = x_img.shape
x_warped: torch.tensor = K.geometry.warp_affine(x_img, M, dsize=(h, w))


Rotate a batch of images#

x_batch = x_img.repeat(16, 1, 1, 1)
x_rot = K.geometry.rotate(x_batch, torch.linspace(0., 360., 16))

imshow(x_rot, (16, 16))