Is it possible to use non-pytoch augmentation in transform.compose

Question

I am working on a data classification problem that takes images as an input in Pytorch. I would like the use the imgaug library, but unfortunatly I keep on getting errors. Here is my code.

#import necessary libraries
from torch import nn
from torchvision import models
import imgaug as ia
import imgaug.augmenters as iaa
from torchvision import datasets
from torch.utils.data.dataloader import DataLoader
from torchvision import transforms
from torch import optim
import numpy as np
from PIL import Image
import glob
from matplotlib import image

#preprocess images
#create data transformers
seq = iaa.Sequential([iaa.Sometimes(0.5,iaa.GaussianBlur(sigma=(0,3.0))),
                      iaa.Sometimes(0.5,iaa.LinearContrast((0.75,1.5))),
                      iaa.AdditiveGaussianNoise(loc=0,scale=(0.0,0.05*255),per_channel=0.5),
                      iaa.Sometimes(0.5,iaa.Affine(
        scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
        translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},
        rotate=(-25, 25),
        shear=(-8, 8)))],random_order=True)
        

 

train_transformation = transforms.Compose([transforms.RandomResizedCrop(300),
                                           seq,
                                           transforms.ToTensor()])

train_data = datasets.ImageFolder(root = 'train')
train_loader = DataLoader(train_data,shuffle = True,batch_size = 32,num_workers = 0)
train_iter = iter(train_loader)
train_iter.next()

----
TypeError                                 Traceback (most recent call last)
 in 
     20 train_loader = DataLoader(train_data,shuffle = True,batch_size = 32,num_workers = 0)
     21 train_iter = iter(train_loader)
---> 22 train_iter.next()

TypeError(default_collate_err_msg_format.format(elem_type))

TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found 

I am aware that the input to the imgaug transformer must be a numpy array, but I am not sure how to incorporate that into my transform.compose (if I can at all that is.). When the imgaug seq is not in the transform.compose it works properly.

Thank you for the help!

Answer 1

Looking at the documentation of transforms in pytorch gives us a hint of how to do it: https://pytorch.org/docs/stable/torchvision/transforms.html#generic-transforms

I would try something like:

train_transformation = transforms.Compose([transforms.RandomResizedCrop(300),
                                           transforms.Lambda(lambda x: seq(x)),
                                           transforms.ToTensor()])