Image and Keypoints augmentations

In this tutorial we leverage kornia.augmentation.AugmentationSequential to apply augmentations to image and transform reusing the applied geometric transformation to a set of associated keypoints. This is useful for detection networks or geometric problems.

Install and get data

%%capture
!pip install kornia
!pip install kornia-rs matplotlib

import io

import requests


def download_image(url: str, filename: str = "") -> str:
    filename = url.split("/")[-1] if len(filename) == 0 else filename
    # Download
    bytesio = io.BytesIO(requests.get(url).content)
    # Save file
    with open(filename, "wb") as outfile:
        outfile.write(bytesio.getbuffer())

    return filename


download_image("https://github.com/kornia/data/raw/main/arturito.jpg")

'arturito.jpg'

import kornia as K
import torch
from matplotlib import pyplot as plt

img = K.io.load_image("arturito.jpg", K.io.ImageLoadType.RGB32)
img = img[None]  # 1xCxHxW / fp32 / [0, 1]
print(img.shape)

torch.Size([1, 3, 144, 256])

Draw points and show image

coords = torch.tensor([[[125, 40.0], [185.0, 75.0]]])  # BxNx2 [x,y]

fig, ax = plt.subplots()

ax.add_patch(plt.Circle((coords[0, 0, 0], coords[0, 0, 1]), color="r"))
ax.add_patch(plt.Circle((coords[0, 1, 0], coords[0, 1, 1]), color="r"))

ax.imshow(K.tensor_to_image(img))

Resize points and show

resize_op = K.augmentation.AugmentationSequential(
    K.augmentation.Resize((100, 200), antialias=True), data_keys=["input", "keypoints"]
)

print(resize_op.transform_matrix)

img_resize, coords_resize = resize_op(img, coords)


fig, ax = plt.subplots()

ax.add_patch(plt.Circle((coords_resize[0, 0, 0], coords_resize[0, 0, 1]), color="r"))
ax.add_patch(plt.Circle((coords_resize[0, 1, 0], coords_resize[0, 1, 1]), color="r"))

ax.imshow(K.tensor_to_image(img_resize))

None

Crop image and points

crop_op = K.augmentation.AugmentationSequential(K.augmentation.CenterCrop((100, 200)), data_keys=["input", "keypoints"])
print(crop_op.transform_matrix)

img_resize, coords_resize = crop_op(img, coords)


fig, ax = plt.subplots()

ax.add_patch(plt.Circle((coords_resize[0, 0, 0], coords_resize[0, 0, 1]), color="r"))
ax.add_patch(plt.Circle((coords_resize[0, 1, 0], coords_resize[0, 1, 1]), color="r"))

ax.imshow(K.tensor_to_image(img_resize))

None