from fastai.vision.all import *

Creating your own Transform

Creating your own Transform is a lot easier than you might think.

Every time you pass a labeled function to a datablock API or to ImageDataLoaders.from_name_func, you create a Transform.

At its base, a Transform is just a fancy function. To get started with easily adding transforms, we’ll implement one that wraps a data augmentation from the albumentations library

!pip install albumentations

Then, it’ll be easier to see the results of the transform on a color image bigger than the MNIST from before, let’s load something from the PETS dataset:

source = untar_data(URLs.PETS)
items = get_image_files(source/"images")

We can open it using PILImage.create:

img = PILImage.create(items[0])
img

Let’s show how to wrap a single transform:

from albumentations import ShiftScaleRotate

ALbumentations transforms work on numpy images, so we just need to convert our PILImage to a numpy array:

aug = ShiftScaleRotate(p=1)
def aug_tfm(img): 
    np_img = np.array(img)
    aug_img = aug(image=np_img)['image']
    return PILImage.create(aug_img)

aug_tfm(img)

We can then pass the function each time a Transform is expected, the fastai library will then automatically do the conversion.

tfm = Transform(aug_tfm)

If you have a state in your transform, you’ll want to create a subclass of Transform.

class AlbumentationsTransform(Transform):
    def __init__(self, aug): self.aug = aug
    def encodes(self, img: PILImage):
        aug_img = self.aug(image=np.array(img))['image']
        return PILImage.create(aug_img)