Artificial Intelligence/Computer Vision

[논문리뷰] CRAFT: Character Region Awareness for Text Detection

oranz.23 2021. 3. 21. 06:15

개요

    이 논문은 Scene Text Detection Task에서 bounding box가 직사각형 모양이 아닌 경우들에 있어서 기존의 모델들이 가지고 있었던 한계점을 돌파하는 데에 그 주안점을 두고 있습니다. 바로 word level이 아니라 character level로 detect한 후 character들 간의 관계를 찾는 방식인데요, 개인적으로 주어진 데이터셋에 문제해결방식을 맞추는 것이 아니라 문제해결방식에 데이터셋을 맞추는 인상을 받아 놀라웠습니다. 그럼 어떻게 ClovaAI가 풀어나갔는지 간단히 살펴보도록 하겠습니다. 이 포스팅에서 첨부하는 이미지 및 코드는 아래 링크들로부터 가져왔음을 밝힙니다.

 

코드: https://github.com/clovaai/CRAFT-pytorch

논문: https://arxiv.org/pdf/1904.01941.pdf

 

Ground Truth Label Generation

   

합성 이미지 라벨링

    먼저, 학습을 위해 합성 이미지에서 라벨링하는 방법입니다. real world dataset들과 다르게 만들어진 이미지들은 연구자들의 의도에 맞게 character 단위로 annotation이 제공됩니다. 그렇지만 이렇게 만들어진 데이터셋 역시 2가지의 score를 위한 라벨링 과정을 거쳐야 합니다. 그 score들은 바로 region score affinity score이며, 각각 픽셀이 character의 중앙에 있을 확률과 픽셀이 인접한 character들의 중앙에 있을 확률을 나타냅니다.

    이제는 해당 score들을 주어진 box에 표시하는 일이 남았습니다. 합성 이미지들인만큼 character box들은 기본적으로 제공되며, affinity box들을 위 그림에서 보이는 바와 같이 각 character box에서 대각선을 그어 만들어진 위아래 삼각형의 중심을 이음으로써 나타낼 수 있습니다. 그리고 상하좌우 대칭인 기본 2D Gaussian Map을 주어진 box에 perspective transform한 후 적용하면 synthetic image들에 대한 ground truth label들이 heatmap 모양으로 완성됩니다.

    

Weakly-Supervised Learning

    synthetic dataset은 character level의 annotation이 있었지만, real world dataset에는 word level까지밖에 존재하지 않습니다. 이를 극복하기 위해 weakly-supervised learning 방식을 사용합니다. weakly-supervised learning이란 Stanford 대학의 연구팀에서 제시된 개념으로, 간단히 말하자면 training data를 프로그래밍을 통해 만들어내는 방식입니다.

 

character splitting

    위와 같이 주어진 word 부분을 crop해내고, 학습된 모델을 통해 word마다 character별로 region score를 예측합니다. 이후에는 watershed algorithm을 사용해 character region을 뽑아내어 이를 바탕으로 character level bounding box가 생성되고 원래 이미지에 반영됩니다. 이러한 일련의 과정을 거쳐 기존의 word-level annotation까지만 되어있던 dataset이 character-level의 정보까지 담게 되었습니다. 이후에는 위의 synthetic dataset을 처리한 것과 마찬가지의 방법으로 ground truth label을 만들어주면 됩니다. 

    그러나 weakly-supervised learning을 통해 탄생한 pseudo ground truth 역시 모델을 통한 예측이므로, 100% 신뢰할 수는 없습니다. 때문에 각각의 pseudo ground truth에 대해 얼마나 신뢰할 수 있는지 계산할 필요가 있습니다.

 

    $l(w)$는 원래 annotation에서의 word 길이를, $l ^{c}(w)$는 pseudo ground truth에서 제시된 해당 word의 길이를, $s _{conf}(w)$는 이를 바탕으로 계산된 해당 word의 confidence score를 나타냅니다. 만약 제시된 character들의 개수가 원래 word의 길이와 같다면 $s _{conf}(w)$는 1이 될 것이지만, 그것보다 적은 수를 찾아냈다면 $s _{conf}(w)$는 그보다 작아집니다. $s _{conf}(w)$가 0.5보다 작을 경우, 수치가 너무 낮으므로 word의 bounding box를 word의 길이만큼 균등하게 나누어서 각 box에 그 순서에 해당하는 character가 있을 것이라고 예측하고 학습합니다. 이 때의 confidence score는 0.5로 주어집니다.

    픽셀 단위로 confidence score를 나타내는 $S _{c}(p)$는 word bounding box에 속할 경우 $s _{conf}(w)$를, 그렇지 않을 경우에는 1을 넣어줌으로써 완성됩니다. 이렇게 완성된 confidence map을 이용해 다음과 같은 loss function으로 학습합니다. synthetic dataset은 bounding box가 character level로 주어지므로 $S _{c}(p)$가 항상 1입니다.

 

    $S _{r}(p)$와 $S _{a}(p)$는 각각 모델에서 예측된 region score와 affinity score를,  $S _{r}^{*}(p)$와 $S _{a}^{*}(p)$는 각각 pseudo ground truth에서의 region score와 affinity score를 나타냅니다. 즉, 각 score들의 ground truth와 예측값 사이의 L2 Error를 구해서 더해준 뒤 confidence score를 곱해서 계산해주는 것입니다.

 

Inference

polygon generation

    추론 과정은 다음과 같습니다. 우선, 이미지 사이즈의 binary map $M$의 값을 모두 0으로 초기화하고, $M(p)$를 예측된 region score와 affinity score에 대해 정해놓은 threshold를 넘긴다면 1로 설정합니다. $M$에 대하여 Connected Component Labeling을 수행합니다. 그리고 모든 라벨에 대응하는 connected component들을 포함하는 최소의 직사각형을 그려 QuadBox를 얻게 됩니다.

   추가적으로, scanning direction을 따라 character별 local maxima로 위 그림에서의 blue arrow들을 그리고, 그 중 가장 긴 것으로 길이를 통일해줍니다. 그리고 각 arrow들의 중간점을 이어 yellow line을 그립니다. 그리고 blue arrow들을 yellow line에 대해 수직이 되도록 회전시켜 red arrow로 만들고, 가장 바깥쪽에 있는 red arrow들을 yellow line을 따라 움직여 text region을 모두 포함할 수 있도록 한 후, 각 red arrow들의 양 끝 점인 control point들을 이으면 text polygon이 완성됩니다.

    Inference 부분의 코드: https://github.com/clovaai/CRAFT-pytorch/blob/master/craft_utils.py