[Deskew for ocr] Rotation correction

2022. 2. 15. 23:50

배경

이번 포스팅에서는 ocr 성능을 높이기 위해 간단한 전처리를 통해 회전된 문서를 올바르게 돌려 주는 문제를 풀어보고자 한다.

OCR 은 optical character recognition의 약자로 hand writing, 인쇄된 문서 등을 카메라로 찍거나 스캔했을때 그 문서 내의 글자를 인식해 문서를 전산화 할때 자주 사용되는 기술이다.

이때 ocr의 성능을 떨어뜨리는 문제 중 하나는 입력으로 들어오는 문서가 회전되어있는 경우 이다.

문서가 회전된 상태로 스캔되면 각 문자 자체는 인식이 할수 있지만 문자를 단어로 머지(merge)하는 과정이나 숫자의 연속인 여권번호, 운전면허 번호, 주민등록 번호 등 긴 문자의 경우 하나의 시퀀스로 머지 되지 않아 잘못된 패턴으로 인식 되는 경우가 종종 발생해 최종 인식률이 떨어질 수 있다.

이런 문제를 해결하기 위해 인식 네트워크 자체가 문서 얼라인먼트 기능을 가지도록 설계할 수도 있고, hand craft feature extraction을 기가 막히게 설계 하고 이를 이용한 homography 계산 알고리즘을 설계하는 방식도 있겠으나!!!

여기서는 간단한 네트워크 설계하고 학습해 rotation correction하는 방식을 시도해 본다.

* 이 포스팅에 사용된 문서는 공공문서 양식 중 하나이고 이런 DB는 'aihub -> 음성/자연어->공공 행정문서 OCR'에서 아주 손쉽게 구할수 있다.

이 포스트에서 개발한 코드는 이 링크에 있다.

code :https://github.com/pajamacoders/ocrDeskew

목표

Fig 1. 왼쪽은 회전된 문서, 오른쪽은 회전 없이 정상적으로 스캔된 문서

이 포스팅의 목적을 명확히 나타내는 그림이 Fig 1. 이다. Fig 1 의 왼쪽은 회전된 문서를 보여준다. 우리의 목표는 이렇게 회전된 문서를 오른 쪽과 같이 회전이 없는 상태로 만드는 것이다.

문제 정의

나는 이 포스팅의 목표인 rotation correction 문제를 image orientation prediction 문제로 정의 했다.

사고의 흐름은 입력 이미지로 부터 이미지의 회전 정도(orientation)를 구하면 그 회전의 크기만큼 반대 방향으로 회전을 시켜 줌으로서 이미지를 정상적으로 만들수 있기 때문이다.

목표 추정 범위는 -30~30degree 이내의 회전으로 정했다.

회전의 크기를 일정 step으로 양자화 해서 각 구간에 class를 부과해 classification으로 해결 할 수도 있을것 같지만

회전 크기를 degree로 바로 추정하는 regression 문제로 정의 하고 풀고자 한다.

개발 환경

개발 환경은 ngc repo에서 아래 이미지를 다운 받았다.

docker image : nvcr.io/nvidia/pytorch:22.01-py3

train metric tracking: mlflow

전처리

전처리는 과정에서는 크게 아래 4가지를 수행했다.

Color conversion(BGR 2 GRAY)
Resize
Rotation (이 과정은 target GT를 생성하는 과정이다.)
Normalization

이때 resize시에 특정한 문제가 발생했고 이를 해결하기 위해 약간의 꼼수(=테크닉)가 사용되어 그 부분에 대해 잠시 설명 하고자 한다.

입력 이미지가 대부분 A4등의 문서를 스캔한 것이기 때문에 풀고자하는 문제에 비해 불필요 하게 크다고 생각했고 입력 이미지의 크기를 512x512로 정의 했다.

하지만 입력 중 어떤 이미지들은 3508x2480(width x height)의 크기를 지녔고 이걸 (512x512)로 리사이즈 하니 다음과 같은 이미지가 생성되었다.

예제의 이미지는 그나마 좀 나은 편인데 심한 것들은 의미있는 글자나 선등이 sampling이 거의 안되어 노이즈 처럼 점만 뿌려져 있는 것처럼 보이기도 한다.

문제는 이러한 텍스트 문서의 경우 배경에 비해 정보를 담고 있는 글자나 테이블을 이루는 선등의 면적이 매우 작아 리사즈 할때 텍스트나 선 부분에서 샘플링되는 양이 매우 작기 때문이다.

위에 서 말했듯이 나는 원인이 정보를 포함하는 텍스트나 선이 차지하는 면적이 작기 때문으로 보았기 때문에 리사이즈 하기 전에 이 정보를 포함하는 부분의 면적을 키우기로 했다.

가장 먼저 떠오른 것은 dilate 연산이다.

dilate을 적용하기 위해 이미지의 색을 반전 시켜 준다. 색 반전의 이유는 입력 이미지는 흰색의 배경을 가지는데 기본적으로 dilate연산은 intensity가 큰 부분은 확장하는 효과를 가지기 때문에 내가 면적을 키우고자 하는 텍스트와 선 등 정보를 가진 pixel이 큰intensity 가지고 background가 낮은 intensity를 가지게 하기 위해서 이다.

아래 코드와 같이 이미지를 읽고 gray scale로 색변환 한 후 background와 foreground의 색상을 반전 시켰다.

    img = cv2.imread(self.img_pathes[index])
    # transform
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    if self.inverse_color:
        img = 255-img # inverse image to apply dilate
    res_dict = {'img':img, 'imgpath':self.img_pathes[index]}

다음으로 적용한 테크닉은 3508 x 2480 을 곧바로 512x512로 리사이즈 하는게 아니라

dilate ->1/2 resize -> 1/2 resize-> resize to 512x512 로 리사이즈 스탭을 진행해 orientation 추출에 필요한 정보를 최대한 많이 보존 하게 하는 것이다.

이 부분의 코드는 다음과 같다.

class Resize(object):
    def __init__(self, scale=4):
        assert (scale!=0) and (scale&(scale-1))==0, 'scale must be power of 2'
        self.iter = np.log2(scale).astype(int)
    
    def __call__(self, inp):
        img = inp['img']
        h,w = img.shape
        inp['org_height']=h
        inp['org_width']=w
        k = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
        img = cv2.dilate(img, k)
        for i in range(self.iter):
            h, w = h//2, w//2
            img=cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
            h, w = img.shape
        inp['img']=cv2.resize(img, (512, 512), interpolation=cv2.INTER_CUBIC)
        return inp

이렇게 만들어진 이미지는 아래 Fig 3과 같다. Fig 2와 텍스트와 선분의 정보 손실 정도를 비교해 보자.

리사이즈 방식에 따라 최종 목적인 regression 성능이 얼마나 달라지는지 ablation study를 진행해 보지는 않았지만 정보의 손실의 최소화 하고자 하는 측면에서는 이게 맞지 않나 싶다.

또 다른 방식으로는 텍스트의 뭉개짐을 무시하고 높은 intensity를 가지는 부분을 최대화 하고자 한다면 아래와 같이 입력을 변환 하는 방식도 가능 할 것이다.

Fig 4에서는 글자를 알아 보기는 힘들지만 이미지의 orientation을 결정하는데에 foreground 정보의 양이 중요할 경우 유용 할 것이다. Fig 4는 위 Resize 클래스 에서 cv2.dilate() 함수를 for 문 안에 cv2.resize함수 호출 바로 앞으로 옮긴 경우 만들어 지는 이미지 이다.

GT 생성

위 이미지 전처리 과정에서 Rotation 과정은 아래와 같은 이유에서 추가되었다.

이미지가 회전된 정도를 GT로 만들어 놓은 경우는 거의 없기 때문에 정상적인 이미지를 opencv를 이용해 랜덤으로 -30~30degree 사이 각도로 회전시켜 회전된 정도를 GT로 사용 한다.

아래코드는 resize된 이미지를 임의의 각도로 회전하고 회전한 각도를 target GT로 저장하는 코드이다.

class RandomRotation(object):
    def __init__(self,degree):
        self.variant = eval(degree) if isinstance(degree, str) else degree

    def __call__(self, inp):
        deg = np.random.uniform(-self.variant, self.variant)
        img = inp['img']
        h,w= img.shape
        matrix = cv2.getRotationMatrix2D((w/2, h/2), deg, 1)
        dst = cv2.warpAffine(img, matrix, (w, h),borderValue=0)
        inp['img'] = dst
        inp['degree'] = deg
        return inp

모델 설계

입력의 회전된 각도를 추정하는 아주 가벼운 네트워크를 아래와 같이 구성했다.

import torch
import torch.nn as nn


class ConvBnRelu(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, dilation=1,
                 groups=1):
        super(ConvBnRelu, self).__init__()
        self.conv_bn_relu = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding, dilation, groups,
                      False),
            nn.BatchNorm2d(out_channel),
            nn.ReLU(True))

    def forward(self, x):
        return self.conv_bn_relu(x)

class DeskewNet(nn.Module):
    def __init__(self, pretrained=None):
        super(DeskewNet, self).__init__()
        k=3
        self.backbone = nn.Sequential(
            ConvBnRelu(1,8,k,padding=k//2),
            nn.MaxPool2d(2,2), #256x256
            ConvBnRelu(8,16,k,padding=k//2),
            nn.MaxPool2d(2,2), #128x128
            ConvBnRelu(16,32,k,padding=k//2),
            nn.MaxPool2d(2,2), #64x64
            ConvBnRelu(32,64,k,padding=k//2),
            nn.MaxPool2d(2,2), #32x32
            ConvBnRelu(64,64,k,padding=k//2),
            nn.MaxPool2d(2,2), #16x16
        )
        self.avgpool = nn.AvgPool2d((16,16))
        self.fc = nn.Sequential(nn.Linear(64,64),nn.Linear(64,1))
        self.__init_weight()
        if pretrained:
            self.load_weight(pretrained)

    def forward(self, x):
        out = self.backbone(x)
        out = self.avgpool(out)
        out = self.fc(out.squeeze())
        return out

    def __init_weight(self):
        for layer in self.modules():
            if isinstance(layer, nn.Conv2d):
                nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
            else: 
                pass

학습

학습 환경은 아래와 같이 구성했다.

optimizer: Adam

lr_scheduler: CosineAnnealingLR -> 초기 lr =0.001, eta_min 1e-6, T_Max=300

loss: MSELoss

batch: 128

아래는 학습 코드 메인 문이다. train, valid 함수 구현은 뻔하니 생략한다.

if __name__ == "__main__":
    args = parse_args()
    with open(args.config, 'r') as f:
        cfg = json.load(f)
        cfg['config_file']=args.config
        if args.run_name:
            cfg['mllogger_cfg']['run_name']=args.run_name
    
    tr = build_transformer(cfg['transform_cfg'])
    train_loader, valid_loader = build_dataloader(**cfg['dataset_cfg'], augment_fn=tr)

    logger.info('create model')
    model = build_model(**cfg['model_cfg'])#torch.hub.load('pytorch/vision:v0.10.0', cfg['model_cfg']['type'])
    model.cuda()
    logger.info('create loss function')
    fn_loss = torch.nn.MSELoss()

    logger.info('create optimizer')
    opt=torch.optim.Adam(model.parameters(), **cfg['optimizer_cfg']['args'])
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt,**cfg['lr_scheduler_cfg']['args'])

    max_epoch = cfg['train_cfg']['max_epoch']
    valid_ecpoh = cfg['train_cfg']['validation_every_n_epoch']
    logger.info(f'max_epoch :{max_epoch}')
    logger.info('set mlflow tracking')
    mltracker = MLLogger(cfg, logger)
    for step in range(max_epoch):
        train(model, train_loader, fn_loss, opt, mltracker, step)
        if (step+1)%valid_ecpoh==0:
            valid(model, valid_loader, fn_loss,  mltracker, step)
        lr_scheduler.step()

결과

아래는 train, validation loss를 mlflow tracking 으로 추적한 결과 이다. loss가 초반에 급격히 떨어지고 수렴은 하지만

validation이 중간 중간 불안정 하게 튀는 모습을 볼 수 있다. 이건 해결을 좀 해야겠다.

아래 Fig 6.의 왼쪽은 회전왼 입력 이미지이고 오른쪽 그림은 orientation을 추정해 correction 한 이미지 이다. 정확하진 않지만 생태는 많이 나아 졌다. 좀더 정확하게 correction 되도록 data검증 및 네트우크 수정을 해봐야겠다.

끝.

'Deeplearning > toyproject' 카테고리의 다른 글

[head pose] head pose 를 이용한 gaze estimation 및 불확실성(uncertainty) 추정 (0)	2022.06.28
[Deskew for ocr] Rotation correction v2 (0)	2022.03.06