How to rotate batch of images by a batch of angles?

vision
1

I have a Tensor of images of size [B, C, H, W] and I want to rotate each of them by a specific angle from the Tensor of angles of size [B].

For example, I have Tensor of images of size [2, 1, 3, 3] and angles of size [2]:

t = torch.Tensor([
    [[[0, 0, 0],
      [2, 0, 1],
      [0, 0, 0]]],
    [[[0, 2, 0],
      [0, 0, 0],
      [0, 1, 0]]]])

r = torch.Tensor([-180, -90])

result should be:

torch.Tensor([
    [[[0, 0, 0],
      [1, 0, 2],
      [0, 0, 0]]],
    [[[0, 0, 0],
      [1, 0, 2],
      [0, 0, 0]]]])

How can I do it?

2

I have already solved this problem but i want to share my solution with community

def rotate(x: torch.Tensor, degree: torch.Tensor) -> torch.Tensor:
    """
    Rotate batch of images [B, C, W, H] by a specific angles [B] 
    (counter clockwise)

    Parameters
    ----------
    x : torch.Tensor
      batch of images
    angle : torch.Tensor
      batch of angles

    Returns
    -------
    torch.Tensor
        batch of rotated images
    """
    angle = degree / 180 * torch.pi
    s = torch.sin(angle)
    c = torch.cos(angle)
    rot_mat = torch.stack((torch.stack([c, -s], dim=1),
                           torch.stack([s, c], dim=1)), dim=1)
    zeros = torch.zeros(rot_mat.size(0), 2).unsqueeze(2)
    aff_mat = torch.cat((rot_mat, zeros), 2)
    grid = F.affine_grid(aff_mat, x.size(), align_corners=True)
    x = F.grid_sample(x, grid, align_corners=True)
    return x
1 Like