Translate, rotate, center crop and resize a batch of images

billiout: <p>I have a batch of images with shape <code>[B, 3, H, W]</code>. 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.</p>
<p>Currently I’m using the following code with <code>torchvision</code> functions <a href="https://pytorch.org/vision/stable/generated/torchvision.transforms.functional.affine.html#torchvision.transforms.functional.affine" rel="noopener nofollow ugc"><code>affine</code></a>, <a href="https://pytorch.org/vision/stable/generated/torchvision.transforms.functional.rotate.html#torchvision.transforms.functional.rotate" rel="noopener nofollow ugc"><code>rotate</code></a>, <a href="https://pytorch.org/vision/stable/generated/torchvision.transforms.functional.center_crop.html#torchvision.transforms.functional.center_crop" rel="noopener nofollow ugc"><code>center_crop</code></a> and <a href="https://pytorch.org/vision/stable/generated/torchvision.transforms.functional.resize.html#torchvision.transforms.functional.resize" rel="noopener nofollow ugc"><code>resize</code></a> but it’s obviously not efficient since it processes each image separately. Any idea how to run these transformations in batch mode?</p>
<pre><code class="lang-auto"># 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)
</code></pre>
<p>Update:</p>
<p>Oddly enough, using the transformation functions from <code>kornia</code> is 3x slower than the above code.</p>
<pre><code class="lang-auto"># 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))
</code></pre>

J_Johnson: <p>Pytorch has image transforms which can be found here:</p>
<p><a href="https://pytorch.org/vision/stable/transforms.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/vision/stable/transforms.html</a></p>
<p>I will implement an example using those modules below. This gives an example based on all of the images provided being uniform in size. If they are not uniform in size, then you will need to iterate through them first to resize them accordingly.</p>
<pre><code class="lang-auto">import torch
from torchvision.transforms import RandomAffine, RandomRotation, CenterCrop, Resize

batch_size = 128
channels = 3
H = 224
W = 224

images = torch.rand((batch_size, channels, H, W))

images = images.to('cuda:0')

affine = RandomAffine(degrees=0., translate=(0.1,0.1))
rotate = RandomRotation(9)
center = CenterCrop((24,24))
resize = Resize((32,32))

images = affine(images)
images = rotate(images)
images = center(images)
images = resize(images)

print(images.size())
</code></pre>
<p>The speed advantage with the above is twofold. One, you will be processing the images asynchronously. Two, you’ll be able to easily send the images to a GPU for processing(i.e. <code>.to('cuda:0')</code>.</p>
<p>Finally, to address the point of your per image translation. This likely won’t yield any benefit over allowing the torchvision module to provide random translations for each image. In the above, you simply specify the maximum parameters and let it sort out the rest.</p>

edgarriba: <p>You can find all this implementations and more in kornia that certainly work for batches and with automatic differentiation</p>
<p><a href="https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform" class="onebox" target="_blank" rel="noopener nofollow ugc">https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform</a></p>

billiout: <p>I combined all transformations (translation, rotation, cropping, and resizing) into a single affine transformation and then applied it using <code>F.affine_grid</code> and <code>F.grid_sample</code>.</p>

edgarriba: <p><a class="mention" href="/u/billiout">@billiout</a> you mean <a href="https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform.get_affine_matrix2d" rel="noopener nofollow ugc">get_affine_matrix2d</a> ?</p>

billiout: <p><a class="mention" href="/u/edgarriba">@edgarriba</a> <code>get_affine_matrix2d</code> is not returning the affine matrix I need. I want to first translate and then rotate whereas <code>get_affine_matrix2d</code> is building an affine matrix that first rotates and then translates. In addition, I want my final affine matrix to be chained with the cropping and resizing operations so that I can avoid building an affined grid on my initial high-resolution images (which is expensive to generate) and instead apply all transformations on the final low-resolution resized space. I believe <a href="https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform.crop_and_resize" rel="noopener nofollow ugc">crop_and_resize</a> is building an affine matrix internally that can be used for this purpose.</p>
<p>Once you have the affine matrix created, you can use <a href="https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform.warp_affine" rel="noopener nofollow ugc">warp_affine</a> or call directly the <code>F.affine_grid</code> and <code>F.grid_sample</code> operations on a normalized and inverted version of the affine matrix.</p>

edgarriba: <p>the “crop” matrix is implemented via <a href="https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform.get_perspective_transform" rel="noopener nofollow ugc">get_perspective_transform</a>. We  noticed at some point that crop and resize was faster using slicing. I think in the augmentations module, you can choose both variants.</p>
<p>If you noticed the affine matrix is a composition of the others (except for translation, I’ll make a quick pr adding get_translation_matrix).</p>
<p>One can easily compose and chain transform based on the requerimients.</p>