[논문 리뷰] Data Augmentation for Scene Text Recognition

텍스트 인식에 포커스가 맞춰진 augmentation이 있을까 싶어 논문을 찾던중 ICCV 2021 학회에서 발표된 STR에서의 Data augmentation 논문이 있어서 정리하려 한다.

 

Abstract 일부

Scene Text Recognition(STR) 모델은 실제 데이터를 사용해서 평가하기 때문에 학습 데이터와 테스트 데이터 분포 간의 불일치는 주로 nosie, artifacts, geometry, structure 등의 영향을 받아서 성능 저하로 이어진다. 본 논문에서는 이를 개선하기 위해 36개의 image augmenation function으로 구성된 STRAug를 소개한다. 각 함수는 자연 장면에서 찾을 수 있거나 카메라 센서에 의해 발생하거나 신호 처리 작업 중 발생하는 이미지 속성을 모방한다.

 

Data Augmentation for STR

위 figure 1은 여러 STR 모델의 baseline 예측 결과와 STRAug를 추가한 경우의 예측 결과를 비교해 보여준다. 체리피킹이긴 하겠지만 그래도 성능이 좋아졌다니 기대된다. 단순 augmentation 만으로 성능을 향상시킬 수 있다는 건 굉장히 나이스하니까.

 

Figure 2는 텍스트 이미지에서 첼린징한 텍스트 형태들을 보여준다. 휘어있거나 그림자가 있거나 회전되었거나 폰트가 특이하거나... 굉장히 많은 경우가 있다.

 

사실 저렇게 다양한 케이스의 텍스트 데이터를 모을 수 있다면 좋겠지만 현실적으로 쉽지 않으니 적절한 augmentation을 사용했을 때 이득을 취할 수 있는 것이다.

 

STRAug

본 논문에서 실제로 적용한 augmentation의 큰 카테고리는 warp, geometry, pattern, noise, blur, weather, camera, process 총 8 가지 이다. 하나씩 사진을 살펴보자. 

 

깔끔한 환경에서 텍스트 인식 작업을 수행하는 경우에는 잘 알 수 없지만, 실제로 환경에서 유입되는 이미지에서 텍스트 검출 후 텍스트 인식을 하게 되면 굉장히 다양한 케이스의 텍스트 이미지를 만나게 된다. 본 논문에서는 굉장히 다양한 케이스를 나누어 효과적인augmentation들을 제안한다. 특히 perspective, 그림자, 블러 등은 굉장히 자주 발생하는 이슈이다. 

 

다만 아쉬운 점은 실제 환경에서의 텍스트 인식은 텍스트 검출 직후 시행되기 때문에 텍스트 box 자체가 정확하지 않은 경우가 많다. 이를 가정해서 텍스트 주변으로 마진을 줄이거나 더 추가하는 등의 augmentation도 있으면 어떨까 싶다. 뜻있는 대학원생이 있다면 실험해보면 좋을 것 같다.

 

 

실험 결과

 

Table 4는 RARE 라는 모델에 다양한 데이터셋 환경에서 앞서 살펴본 8 가지 augmetnation을 각각 적용시켜본 결과이다. 역시나 늘 augmenation이 항상 옳은 것은 아니다. 그래도 전반적으로 많은 케이스의 augmentation이 효과적인 것을 볼 수 있다. Data augmentation은 50% 의 확률로 적용된 것이라 한다.

 

Table 5는 8 가지의 augmentation이 조합된 STRAug와 다른 augmentation 과의 성능 상승 비교를 보여준다. STRAug가 다른 augmentation에 비해 상당히 효과적으로 보인다.

 

---

하지만, 퍼블릭 데이터셋에서의 실험이고 한국어 데이터셋에서는 검증되지 않은 augmentation이다. 또한 실제 환경에서는 accuracy 1~2% 정도의 상승은 그다지 눈여겨 볼만한 차이를 가져다 주지 않을지도 모른다. 그렇지만 텍스트 인식 모델 학습을 위해 충분히 고려해볼만한 data augmentation 방법인 것 같다. 텍스트 인식 모델의 목적에 따라 augmentation 종류를 선별해서 사용하는 것도 나쁘지 않아 보인다. 

 

공식 깃허브 레퍼지토리를 참고하면 굉장히 사용하기 편하게 되어있지만, 테스트 코드가 StrAug 의 여러 augmentation 기법을 하나씩 적용해서 결과를 보여주는 것이라 실제 모델 학습에 바로 적용시키기에는 부적절하다.

때문에 원하는 augmentation 유형을 골라서 랜덤하게 augmentation을 적용한 결과를 보여주는 테스트 코드를 추가하여 pull request를 해놓은 상태이다. 적용될지 안될지 몰라서 작성한 코드를 아래에 첨부한다.

 

import argparse
import os
import numpy as np
from PIL import Image
import cv2

from straug.warp import Curve, Distort, Stretch
from straug.geometry import Rotate, Perspective, Shrink, TranslateX, TranslateY
from straug.blur import GaussianBlur, DefocusBlur, MotionBlur, GlassBlur, ZoomBlur
from straug.camera import Contrast, Brightness, JpegCompression, Pixelate
from straug.noise import GaussianNoise, ShotNoise, ImpulseNoise, SpeckleNoise
from straug.pattern import VGrid, HGrid, Grid, RectGrid, EllipseGrid
from straug.process import Posterize, Solarize, Invert, Equalize, AutoContrast, Sharpness, Color
from straug.weather import Fog, Snow, Frost, Rain, Shadow

class Random_StrAug(object):
    def __init__(self, using_aug_types, prob_list = None):
        self.aug_list = []
        if 'warp' in using_aug_types :
            self.aug_list.append([Curve(), Distort(), Stretch()]) 
        if 'geometry' in using_aug_types :
            self.aug_list.append([Rotate(), Perspective(), Shrink(), TranslateX(), TranslateY()]) 
        if 'blur' in using_aug_types :
            self.aug_list.append([GaussianBlur(), DefocusBlur(), MotionBlur(), GlassBlur(), ZoomBlur()]) 
        if 'noise' in using_aug_types :
            self.aug_list.append([GaussianNoise(), ShotNoise(), ImpulseNoise(), SpeckleNoise()]) 
        if 'camera' in using_aug_types :
            self.aug_list.append([Contrast(), Brightness(), JpegCompression(), Pixelate()]) 
        if 'pattern' in using_aug_types :
            self.aug_list.append([VGrid(), HGrid(), Grid(), RectGrid(), EllipseGrid()]) 
        if 'process' in using_aug_types :
            self.aug_list.append([Posterize(), Solarize(), Invert(), Equalize(), AutoContrast(), Sharpness(), Color()]) 
        if 'weather' in using_aug_types :
            self.aug_list.append([Fog(), Snow(), Frost(), Rain(), Shadow()]) 
    
        self.mag_range = np.random.randint(-1, 3)
        if prob_list is None :
            self.prob_list = [0.5] * len(self.aug_list)
        else:
            assert len(self.aug_list) == len(prob_list), "The length of 'prob_list' must be the same as the number of augmentations used."
            self.prob_list = prob_list

    def __call__(self, img):
        for i in range(len(self.aug_list)):
            img = self.aug_list[i][np.random.randint(0, len(self.aug_list[i]))](img, mag = self.mag_range, prob = self.prob_list[i])

        return img
    
if __name__ == '__main__':
    random_StrAug_1 = Random_StrAug(using_aug_types = ['warp', 'geometry', 'blur', 'noise', 'camera', 'pattern', 'process', 'weather'],
                                  prob_list = [0.5, 0.3, 0.3, 0.2, 0.2, 0.1, 0.1, 0.1])
    
    random_StrAug_2 = Random_StrAug(using_aug_types = ['warp', 'pattern', 'process', 'weather'],
                                  prob_list = [0.5, 0.3, 0.2, 0.5])
    
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', default="images/delivery.png", help='Load image file')
    parser.add_argument('--results', default="results", help='Folder for augmented image files')
    parser.add_argument('--gray', action='store_true', help='Convert to grayscale 1st')
    parser.add_argument('--width', default=200, type=int, help='Default image width')
    parser.add_argument('--height', default=64, type=int, help='Default image height')
    parser.add_argument('--seed', default=0, type=int, help='Random number generator seed')
    opt = parser.parse_args()
    os.makedirs(opt.results, exist_ok=True)

    img = Image.open(opt.image)
    img = img.resize((opt.width, opt.height))

    augmented_img_1 = random_StrAug_1(img)
    augmented_img_2 = random_StrAug_2(img)

    # Save images to compare before and after augmentation.
    result = cv2.cvtColor(np.hstack((np.array(img), np.array(augmented_img_1), np.array(augmented_img_2))), cv2.COLOR_RGB2BGR)
    cv2.imwrite(os.path.join(opt.results, opt.image.split('/')[-1].split('.')[0] + '_random_strAug.jpg'), result)