Random rotation for 5D tensor in pytorch?

1

I have images and targets size of BxCxDxHxW. I want to perform random rotation using Tensor. Currently, I was successful to use ndimage.affine_transform by feeding each image with size of DxHxW to the function likes images = ndimage.affine_transform(images, full_rot_mat) where full_rot_mat is a rotation matrix. I have two questions:

  1. How we can use the below code for tensor; for now, I have to convert back to numpy to use them (img_np = images.data.cpu().numpy())?

  2. How can we rotate a batch (like B) ? For now, I have to rotate each image in the batch likes img_np_=img_np[0,:,:,:,0]?

This is the code what I am using

def do_random_transform(images, targets, x_rotation_max_angel_deg, y_rotation_max_angel_deg, z_rotation_max_angel_deg):
    x_rot_mat = np.eye(3, 3)
    if x_rotation_max_angel_deg > 0:
        rot_deg = random.randint(-1 * x_rotation_max_angel_deg, x_rotation_max_angel_deg)
        x_rot_mat[1, 1] = math.cos(math.radians(rot_deg))
        x_rot_mat[2, 2] = math.cos(math.radians(rot_deg))
        x_rot_mat[1, 2] = -1 * math.sin(math.radians(rot_deg))
        x_rot_mat[2, 1] = math.sin(math.radians(rot_deg))

    y_rot_mat = np.eye(3, 3)
    if y_rotation_max_angel_deg > 0:
        rot_deg = random.randint(-1 * y_rotation_max_angel_deg, y_rotation_max_angel_deg)
        y_rot_mat[0, 0] = math.cos(math.radians(rot_deg))
        y_rot_mat[2, 2] = math.cos(math.radians(rot_deg))
        y_rot_mat[2, 0] = -1 * math.sin(math.radians(rot_deg))
        y_rot_mat[0, 2] = math.sin(math.radians(rot_deg))

    z_rot_mat = np.eye(3, 3)
    if z_rotation_max_angel_deg > 0:
        rot_deg = random.randint(-1 * z_rotation_max_angel_deg, z_rotation_max_angel_deg)
        z_rot_mat[0, 0] = math.cos(math.radians(rot_deg))
        z_rot_mat[1, 1] = math.cos(math.radians(rot_deg))
        z_rot_mat[0, 1] = -1 * math.sin(math.radians(rot_deg))
        z_rot_mat[1, 0] = math.sin(math.radians(rot_deg))

    full_rot_mat = np.dot(np.dot(x_rot_mat, y_rot_mat), z_rot_mat)

    images = ndimage.affine_transform(images, full_rot_mat)
    targets = ndimage.affine_transform(targets, full_rot_mat)

    return images, targets

For using

    images, targets = data #This is tensor
    img_np = images.data.cpu().numpy()
    label_np = targets.data.cpu().numpy()
    img_np_=img_np[0,:,:,:,0]
    label_np_=label_np[0,:,:,:,0]
    images_val, label_val = do_random_transform(img_np_, label_np_, 10, 10, 15)