Using imgaug with PyTorch

imgaug_pytorch_sample.py

# Import PyTorch
import torchvision.transforms.functional as TF
from torch.utils.data import Dataset, DataLoader

# Define the augmentations
AUG_TRAIN = aug_pipeline # use our pipeline as train augmentations

# Define the demo dataset
class DogDataset(Dataset):
    '''
    Sample dataset for imgaug demonstration.
    The dataset will consist of just one sample image.
    '''

    def __init__(self, image, augmentations = None):
        self.image = image
        self.augmentations = augmentations # save the augmentations

    def __len__(self):
        return 1 # return 1 as we have only one image

    def __getitem__(self, idx):
        # return the augmented image
        return TF.to_tensor(self.augmentations.augment_image(self.image))
        
# Load the augmented data

# Initialize the dataset, pass the augmentation pipeline as an argument to init function
train_ds = DogDataset(image, augmentations = AUG_TRAIN)

# Initilize the dataloader
trainloader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=0)