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.

%%capture
!pip install kornia
%%capture
!wget https://github.com/kornia/data/raw/main/bennett_aden.png
import cv2
from matplotlib import pyplot as plt
import numpy as np

import torch
import torchvision
import kornia as K
/home/docs/checkouts/readthedocs.org/user_builds/kornia-tutorials/envs/latest/lib/python3.7/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  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)
    plt.figure(figsize=size)
    plt.imshow(out_np)
    plt.axis('off')
    plt.show()
imshow(x_img)
_images/4c34953e466666a7c2a47ff123514c07a5827a1c03f103134c8f8d41a32f964a.png

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))

imshow(x_warped)
_images/c94fcccfa46cab3c13d714a45313c6efca1ad3244891fd3a4ec27aee2901ab62.png

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))
_images/ee6aaad2b9c8fd333a21bee7ccc08f0b18fac29d0ba8d4be7a5139a73a16d21d.png