TrOCR: Transformer-based Optical Character Recognitionwith Pre-trained Models

TL;DR

• 기존 AI OCR 한계점 : 이미지 이해(CNN), 문자 단위 자연어 생성(RNN), 추가 후처리 (언어모델)
• 트랜스포머 구조 TrOCR은 end-to-end로 하나의 모델로 OCR을 위한 이미지 이해, 토큰 단위 자연어 생성의 과정을 수행
  • 이미지 인코더(ViT), 자연어 디코더
  • 초기화를 각각 ViT 스타일 인코더, 자연어 디코더에서 하는 게 좋다고 추천함
• 뿐만 아니라 TrOCR은 합성 데이터에 사전 훈련 후, 커스텀 데이터셋에 미세 조정을 진행하면 되기 때문에
  용도에 맞게 활용이 매우 용이하다는 측면에서 기여하는 바가 크다.
• 성능 측면에서도 (CER) 인쇄본, 손글씨, scene-text 등 다양한 시나리오의 데이터셋에 대해서도 잘함.

---

TrOCR 구조

그림 1. TrOCR의 아키텍처. 일반 데이터셋에 사전 학습된 이미지 트랜스포머를 인코더로, 일반 텍스트 데이터에 사전 학습된 텍스트 트랜스포머를 디코더로 사용하는 인코더-디코더 모델의 구성.

이미지 이해를 위한 ViT 스타일의 인코더와 (BeIT, DeIT), 자연어 생성을 위한 디코더 (RoBERTA)로 구성됨.
이미지 이해는 ViT같이 이미지를 패치 단위로 나누어 각각을 토큰 취급하여 처리.

디코딩은 이미지 feature, 지금까지 생성된 토큰 정보들을 기반으로 하여 다음 토큰을 자기회귀적으로 예측.

훈련 : 정말로 용이한지?

TrOCR은 텍스트 인식 학습만으로 훈련 가능하며, 텍스트-감지(bbox) 정보가 필요하지 않다.
저자는 OCR 모듈의 출력 토큰 시퀀스와 정답(GT) 토큰 시퀀스 간 CrossEntropy 손실로 훈련했다고 설명했다.

사전 훈련 : 합성 데이터셋

 

• 1단계: 대규모 데이터셋 학습
  • 수억 개의 인쇄된 텍스트 라인 이미지로 구성된 합성 데이터셋을 사용
  • 기본적인 텍스트 인식과 시각적 특징 추출 능력을 학습
• 2단계: 특정 태스크 데이터셋 학습
  • 첫 번째 단계에서 학습된 모델을 초기화 상태로 사용
  • 인쇄된 텍스트와 손글씨 텍스트 인식을 위한 두 개의 작은 데이터셋(각각 수백만 개 이미지)으로 학습
  • 장면 텍스트(Scene Text) 인식의 경우 널리 사용되는 합성 텍스트 데이터셋 활용

 

 

미세 조정 : 커스텀 데이터셋 

 

결과 및 고찰

다양한 벤치마크에서 평가해본 결과

---

Reference

[1] Li, M., Lv, T., Chen, J., Cui, L., Lu, Y., Florencio, D., ... & Wei, F. (2023, June). Trocr: Transformer-based optical character recognition with pre-trained models. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 37, No. 11, pp. 13094-13102).