Hi all,
I implemented some affine transforms for pytorch – specifically, Rotation(), Translation(), Shear(), and Zoom(), and an over-arching Affine() transform which can perform all of those transforms while only using one interpolation.
Right now it’s not maximally efficient because i cast to and from numpy… Eventually I’ll implement this all in torch and then it can be performed efficiently on the gpu.
Here is a link to the gist:

  
      gist.github.com
  
  
    https://gist.github.com/ncullen93/425ca642955f73452ebc097b3b46c493 1.1k
affine_transforms.py
"""
Affine transforms implemented on torch tensors, and
only requiring one interpolation

Included:
- Affine()
- AffineCompose()
- Rotation()
- Translation()
- Shear()
This file has been truncated. show original




  
  
    
    
  
  


You’ll see that the transforms take in an option y argument which allows you to perform the same affine transform on both the input and target image… You’ll need a special dataset class for this, which I’ve provided as a starting point in the following gist for both in-memory (TensorDataset) and out-of-memory (FolderDataset) data… it basically just required adding a co_transform argument:

  
      gist.github.com
  
  
    https://gist.github.com/ncullen93/14b458cb4bd237bab2a41a185f710808 242
co_classes.py
"""
Custom datasets from both in-memory and out-of-memory data
"""

import torch.utils.data as data

from PIL import Image
import os
import os.path

This file has been truncated. show original




  
  
    
    
  
  


I’ve spot-checked these transforms so I’m confident they aren’t blatantly wrong, but if you come across a bug then add a comment in the gist.
Short example of the easiest use case
import torch
affine_transform = Affine(rotation_range=30, translation_range=(0.2,0.2), zoom_range=(0.8,1.2))
x = torch.ones(10,3,50,50)
y = torch.ones(10,3,50,50)
data  = TensorDataset(x, y, co_transform=affine_transform)

x_sample, y_sample = data[0] # the transforms should be applied to both x and y