365. Applying Data Augmentation Using Albumentation

When applying data augmentation for computer vision models, it is important to apply the same augmentation in the same order for both the input and target image. Here is one way to implement using Albumentation and Pytorch.

Implementation

Today I’ll only share how to define the augmentation pipeline and the dataset class.

1. Define Data Augmentation Pipeline
import albumentations as A

#Use Compose to define the procedures for augmentation
transform = A.Compose(
    [
    # Write down whatever augmentation you want to apply
        A.Resize(512, 512),
        A.ShiftScaleRotate(shift_limit=0.7, scale_limit=0.4, rotate_limit=20, p=0.5),
        A.RGBShift(r_shift_limit=25, g_shift_limit=25, b_shift_limit=25, p=0.5),
        A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ]
)
2. Define Dataset Class
class MyDataset(Dataset):

    def __init__(self, images_filenames,mask_filenames, transform=None):
        self.images_filenames = images_filenames
        self.mask_filenames = mask_filenames
        self.transform = transform

    def __len__(self):
        return len(self.images_filenames)

    def __getitem__(self, idx):
        #Get Image Path
        image_filename = self.images_filenames[idx]

        #Open Image
        image = cv2.imread(image_filename)

        #Since OpenCV reads images as BGR format, convert to RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        #Get Target Image
        mask_filename = self.mask_filenames[idx]

                #Open Target Image
        mask = cv2.imread(mask_filename,cv2.IMREAD_UNCHANGED)

        #Apply Data Augmentation for both image and target
        if self.transform is not None:
            transformed = self.transform(image=image, mask=mask)
            image = transformed["image"]
            mask = transformed["mask"]

        return image, mask

Reference

Image from Albumentation

Related Posts