I have a batch of images with shape [B, 3, H, W]
. For each image in the batch, I want to translate it by a pixel location different for each image, rotate it by an angle different for each image, center crop it by its own crop size, and finally, resize them to the same size.
Currently I’m using the following code with torchvision
functions affine
, rotate
, center_crop
and resize
but it’s obviously not efficient since it processes each image separately. Any idea how to run these transformations in batch mode?
# images - [B, 3, H, W]
# translation_point - [B, 2]
# angle - [B, 1]
# crop_size - [B, 1] Assume square cropping
transformed_images = []
for i in range(images.shape[0]):
img = affine(images[i], angle=0, scale=1.0, shear=0.0,
translate=((images.shape[3]/2 - translation_point[i, 1].item()),
(images.shape[2]/2 - translation_point[i, 0].item())))
img = rotate(img, angle[i, 0].item(), expand=False)
img = center_crop(img, (crop_size[i].item(), crop_size[i].item()))
img = resize(img, 64)
transformed_images.append(img)
transformed_images = torch.stack(transformed_images, dim=0)
Update:
Oddly enough, using the transformation functions from kornia
is 3x slower than the above code.
# images - [B, 3, H, W]
# translation_point - [B, 2]
# angle - [B, 1]
# crop_size - [B, 1] Assume square cropping
translation_point[..., 0], translation_point[..., 1] = (images.shape[3]/2 - translation_point[..., 1]),\
(images.shape[2]/2 - translation_point[..., 0])
images = kornia.geometry.transform.translate(images, translation_point)
images = kornia.geometry.transform.rotate(images, angle)
crop_bboxes = kornia.geometry.bbox.bbox_generator(images.shape[2]/2 - crop_size/2,
images.shape[3]/2 - crop_size/2,
crop_size, crop_size)
images = kornia.geometry.transform.crop_and_resize(images, crop_bboxes, (64, 64))