How do I create a shear matrix for PyTorch's F.affine_grid & F.grid_sample?

Question

I need to create a shear matrix that is autograd compatible, works on B,C,H,W tensors, and takes input values (possibly generated randomly) for the shear values. How can I generate the shear matrix for this?

import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image


# Load image
def preprocess_simple(image_name, image_size):
    Loader = transforms.Compose([transforms.Resize(image_size), transforms.ToTensor()])
    image = Image.open(image_name).convert('RGB')
    return Loader(image).unsqueeze(0)
    
# Save image   
def deprocess_simple(output_tensor, output_name):
    output_tensor.clamp_(0, 1)
    Image2PIL = transforms.ToPILImage()
    image = Image2PIL(output_tensor.squeeze(0))
    image.save(output_name)


def get_shear_mat(theta):
    ...
    return shear_mat


def shear_img(x, theta, dtype):
    shear_mat = get_shear_mat(theta)
    grid = F.affine_grid(shear_mat , x.size()).type(dtype)
    x = F.grid_sample(x, grid)
    return x


# Shear tensor
test_input = # Test image
shear_values = (3,4) # Example values
sheared_tensor = shear_img(test_input, shear_values)

user_4757715 · Answer

Say m is the shear factor, then theta = atan(1/m) is the shear angle. You can now pick either horizontal shear or vertical shear. Here's how you implement get_shear_mat such that you can pick horizontal shear by setting ax=0 and vertical shear by setting ax=1:

def get_shear_mat(theta, ax=0):
    assert ax in [0, 1]
    m = 1 / torch.tan(torch.tensor(theta))
    if ax == 0: # Horizontal shear
        shear_mat = torch.tensor([[1, m, 0],
                         [0, 1, 0]])
    else: # Vertical shear
        shear_mat = torch.tensor([[1, 0, 0],
                         [m, 1, 0]])
    return shear_mat

Notice that a shear mapping is just a mapping of point (x,y) in the original image to the point (x+my,y) for horizontal shear, and (x,y+mx) for vertical shear. This is exactly what we do here by defining the shear_mat as above.

An optional modification to shear_img to support the operation for a batched input in the first row. Also adding an argument - ax to shear_img to define whether we want a horizontal (ax=0) or vertical(ax=1) shear:

def shear_img(x, ax, theta, dtype):
    shear_mat = get_shear_mat(theta, ax)[None, ...].type(dtype).repeat(x.shape[0], 1, 1)
    grid = F.affine_grid(shear_mat , x.size()).type(dtype)
    x = F.grid_sample(x.type(dtype), grid)
    return x

Let's test this implementation on an image:

# Let im be a 4D tensor of shape BxCxHxW (an image or a batch of images):
dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor # Set type of data
sheared_im = shear_img(im, 0, np.pi/4, dtype) #Horizontal shear by shear angle of pi/4
plt.imshow(sheared_im.squeeze(0).permute(1,2,0)/255)
plt.show()

If im is our dancing cat with a skirt:

Then our plot will be:

If we want a vertical shear:

sheared_im = shear_img(im, 1, np.pi/4, dtype) # Vertical shear by shear angle of pi/4
plt.imshow(sheared_im.squeeze(0).permute(1, 2, 0)/255)
plt.show()

We obtain:

Hooray!

How do I create a shear matrix for PyTorch's F.affine_grid & F.grid_sample?

1 Answers