Augmentation(Albumentation) 적용하기

[bellamy-k]

아래 Augmentation은 LaTeX-OCR 코드에 사용된 Augmentation입니다.

train_transform = alb.Compose(
        [
            alb.Resize(options.input_size.height, options.input_size.width),
            alb.Compose(
                [alb.ShiftScaleRotate(shift_limit=0, scale_limit=(-.15, 0), rotate_limit=1, border_mode=0,
                                      interpolation=3,
                                      value=[255, 255, 255], p=1),
                 alb.GridDistortion(distort_limit=0.1, border_mode=0, interpolation=3, value=[255, 255, 255], p=.5)],
                p=.15),
            # alb.InvertImg(p=.15),
            alb.RGBShift(r_shift_limit=15, g_shift_limit=15,
                         b_shift_limit=15, p=0.3),
            alb.GaussNoise(10, p=.2),
            alb.RandomBrightnessContrast(.05, (-.2, 0), True, p=0.2),
            alb.JpegCompression(95, p=.3),
            alb.ToGray(always_apply=True),
            alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
            # alb.Sharpen()
            ToTensorV2(),
        ]
    )
    test_transform = alb.Compose(
        [
            alb.Resize(options.input_size.height, options.input_size.width),
            alb.ToGray(always_apply=True),
            alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
            # alb.Sharpen()
            ToTensorV2(),
        ]
    )

위 transform을 예시로 albumentation 적용 방법을 알아보면..

1. train_transform과 test_transform의 분리

우선 우리 코드에선 train.py에서 transform을 정의하고 인자로 넘겨주기에 train.py 부분에서 해당 부분을 바꾸어 주어야 합니다.

before

transformed = transforms.Compose(
    [
        # Resize so all images have the same size
        transforms.Resize((options.input_size.height, options.input_size.width)),
        transforms.ToTensor(),
    ]
  )
train_data_loader, validation_data_loader, train_dataset, valid_dataset = dataset_loader(options, transformed) 

after

train_transform = alb.Compose(
        [
            alb.Resize(options.input_size.height, options.input_size.width),
            alb.Compose(
                [alb.ShiftScaleRotate(shift_limit=0, scale_limit=(-.15, 0), rotate_limit=1, border_mode=0,
                                      interpolation=3,
                                      value=[255, 255, 255], p=1),
                 alb.GridDistortion(distort_limit=0.1, border_mode=0, interpolation=3, value=[255, 255, 255], p=.5)],
                p=.15),
            # alb.InvertImg(p=.15),
            alb.RGBShift(r_shift_limit=15, g_shift_limit=15,
                         b_shift_limit=15, p=0.3),
            alb.GaussNoise(10, p=.2),
            alb.RandomBrightnessContrast(.05, (-.2, 0), True, p=0.2),
            alb.JpegCompression(95, p=.3),
            alb.ToGray(always_apply=True),
            alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
            # alb.Sharpen()
            ToTensorV2(),
        ]
    )

test_transform = alb.Compose(
    [
        alb.Resize(options.input_size.height, options.input_size.width),
        alb.ToGray(always_apply=True),
        alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
        # alb.Sharpen()
        ToTensorV2(),
    ]
)

train_data_loader, validation_data_loader, train_dataset, valid_dataset = dataset_loader(options, train_transform, test_transform)    #수정됨

2. dataset_loader 변경

dataset_loader가 두 가지 transform을 받게 만들었으니 수정해 주어야 합니다.
선언부에서 변수 두 개 받게 하고, dataset.py 내부 dataset_loader 함수의 LoadDataset 내부 transform 인자만 두 개로 나누어 주면 됩니다.

dataset.py

before

def dataset_loader(options, transformed):    #수정됨

    .
    .
    .

    # Load data
    train_dataset = LoadDataset(
        train_data, options.data.token_paths, crop=options.data.crop, transform=transformed, rgb=options.data.rgb
    )
    train_data_loader = DataLoader(
        train_dataset,
        batch_size=options.batch_size,
        shuffle=True,
        num_workers=options.num_workers,
        collate_fn=collate_batch,
    )

    valid_dataset = LoadDataset(
        valid_data, options.data.token_paths, crop=options.data.crop, transform=transformed, rgb=options.data.rgb
    )

    .
    .
    .

    return train_data_loader, valid_data_loader, train_dataset, valid_dataset

after

def dataset_loader(options, train_transformed, test_transformed):    #수정됨

    .
    .
    .

    # Load data
    train_dataset = LoadDataset(
        train_data, options.data.token_paths, crop=options.data.crop, transform=train_transformed, rgb=options.data.rgb    #수정됨
    )
    train_data_loader = DataLoader(
        train_dataset,
        batch_size=options.batch_size,
        shuffle=True,
        num_workers=options.num_workers,
        collate_fn=collate_batch,
    )

    valid_dataset = LoadDataset(
        valid_data, options.data.token_paths, crop=options.data.crop, transform=test_transformed, rgb=options.data.rgb    #수정됨
    )

    .
    .
    .

3. LoadDataset Class 변경

Albumentation의 to_gray는 1채널이 아닌 3채널 이미지로 나오고, np.array로 변환해 주어야 오류가 안 나기에 LoadDataset Class 를 다음과 같이 수정해줍니다.

dataset.py 내부 LoadDataset class의 getitem()

** before **

class LoadDataset(Dataset):
.
.
.

    def __getitem__(self, i):
        item = self.data[i]
        image = Image.open(item["path"])
        if self.rgb == 3:
            image = image.convert("RGB")
        elif self.rgb == 2:
            image = image.convert("L")
            image = image.convert("RGB")
        elif self.rgb == 1:
            image = image.convert("L")
        else:
            raise NotImplementedError

        if self.crop:
            # Image needs to be inverted because the bounding box cuts off black pixels,
            # not white ones.
            bounding_box = ImageOps.invert(image).getbbox()
            image = image.crop(bounding_box)

        if self.transform:
            image = self.transform(image)

        return {"path": item["path"], "truth": item["truth"], "image": image}
.
.
.

** after **

    def __getitem__(self, i):
        item = self.data[i]
        image = Image.open(item["path"])
        #self.rgb 는 모두 1이라고 가정(조건문 삭제)

        if self.crop:
            # Image needs to be inverted because the bounding box cuts off black pixels,
            # not white ones.
            bounding_box = ImageOps.invert(image).getbbox()
            image = image.crop(bounding_box)

        image = np.array(image)    #수정됨

        if self.transform:
            image = self.transform(image=image)['image'][:1]    #수정됨

        return {"path": item["path"], "truth": item["truth"], "image": image}

4. inference 를 위한 수정

추후 inference 때 Albumentation transform을 적용하고 싶다면, dataset.py 내부 LoadEvalDataset class의 getitem() 도 동일하게 수정해 주어야 합니다. 또한 inference.py 내부 transfomed 변수 를 수정해 주어야 합니다

[Sooho-Kim]

친절한 갓정환님