PyTorch 2D affine_grid, grid_sample example

grid_example.py

import numpy as np
import cv2
import torch
from torch import nn
import math

a = np.zeros((100, 100, 3), dtype=np.float32)
a[25:75, 25:75] = np.array((1, 1, 1), dtype=np.float32)
cv2.imshow('Source', (a*255).astype(np.uint8))

a = torch.tensor(a, dtype=torch.float32)
a = a.permute(2, 0, 1)
a = a.view(1, *a.shape)

# In this example, we rotate the image 30 degree clockwise
# than translate quarter of width and height to bottom-right
cos = math.cos(math.radians(30))
sin = math.sin(math.radians(30))
grid = torch.zeros((1, 3, 3), dtype=torch.float32)
# [ cos, -sin, trans_x]
# [ sin,  cos, trans_y]
# [   0,    0,       1]
# Note that translations should be normalized to [-1.0, 1.0]
grid[:, 0, 0] = cos
grid[:, 0, 1] = -sin
grid[:, 1, 0] = sin
grid[:, 1, 1] = cos
grid[:, 0, 2] = 0.5 # right
grid[:, 1, 2] = 0.5 # bottom
grid[:, 2, 2] = 1
# Since affine_grid expectes 'target' to 'source' mapping (inverse mapping),
# we have to invert this transform matrix for 'source' to 'target' mapping
grid = torch.inverse(grid)[:, :2]

b = nn.functional.affine_grid(grid, a.shape, align_corners=False)
b = nn.functional.grid_sample(a, b)
b = b.squeeze(dim=0) # (3, 12, 12)
b = b.permute(1, 2, 0)
b = (b.numpy()*255).astype(np.uint8)
cv2.imshow('Target', b)
cv2.waitKey(0)