OCRモデルTrOCRについて調べたまとめ

お疲れ様です。

今月の半ばくらいまで会社の勉強会の関係でAI OCRの手法の1つのTrOCRの調査と実装をしていました。今回はそちらの内容をまとめて備忘録を書いておきます。

論文

要約（from ChatGPT）

TrOCR：Transformerベースの革新的なOCR技術

TrOCRは「Transformer Optical Character Recognition」の略で、Microsoftが開発した、Transformerアーキテクチャに基づく先進的な光学文字認識（OCR）モデルです。従来のOCR技術が主に畳み込みニューラルネットワーク（CNN）を利用していたのに対し、TrOCRは画像認識とテキスト生成の両方にTransformerを用いることで、シンプルかつ効果的なエンドツーエンドのシステムを実現しています。

主な特徴

Transformerのみの構成: TrOCRは、画像処理のためのCNNを必要とせず、画像Transformer（エンコーダ）とテキストTransformer（デコーダ）を組み合わせたシンプルな構成が特徴です。これにより、入力画像特有のバイアスが少なく、高い汎用性を持ちます。

エンドツーエンドのシステム: 画像の入力からテキストの出力までを一つのモデルで完結させるエンドツーエンドのアプローチを採用しています。

事前学習の活用: 大規模な合成データセットで事前にモデルを学習（Pre-training）させ、その後、特定のタスク（手書き文字、印刷文字など）に合わせて微調整（Fine-tuning）することで、非常に高い認識精度を達成します。

高い認識精度: 手書き文字と印刷されたテキストの両方で、従来のモデルを上回る優れた性能を示しています。

仕組み

TrOCRのアーキテクチャは、主に2つの部分から構成されています。

画像エンコーダ（Image Transformer）: 入力された画像を小さなパッチに分割し、それぞれのパッチから特徴を抽出して文脈を捉えます。これは、コンピュータービジョンの分野で高い性能を示すVision Transformer（ViT）と同様のアプローチです。

テキストデコーダ（Text Transformer）: 画像エンコーダが抽出した特徴情報を受け取り、それを基に対応するテキストシーケンスを生成します。

このエンコーダ・デコーダ構造により、画像の内容を理解し、それを自然言語のテキストとして出力することが可能になります。

ざっくりとアーキテクチャをすべてTransformerベースで作成したOCRモデルという感じですね。 EncoderとDecoderの２つのアーキテクチャを用います。それぞれ以下のようなアーキテクチャを使用します。

Encoder -> 画像系Transformer（ViTやDeiTなど）
Decoder -> 自然言語系Transformer（BERTやGPTなど）

arch

ライセンスはMITです。個人的には実際に業務でも実装した経験があります。

公式実装は以下（2つある…？）
github.com github.com

テストコードの実装

確認用のテストコードの実装を記載します。
TrOCRはHuggingFaceから利用可能でMicrosoftが作成したモデルを使用することができます。ただし、これらは英語に特化したモデルで日本語ではそのまま使えません。

使用可能なモデル一覧

モデル名	学習段階	主なデータ	用途	備考
trocr-base-stage1	事前学習	合成文字画像	再学習ベース	Fine-tuning前提
trocr-base-printed	Fine-tuning	印刷体	文書OCR	通常の印刷OCR向け
trocr-base-handwritten	Fine-tuning	手書き文字	手書きOCR	筆記体対応
trocr-base-str	Fine-tuning	自然画像文字（STR）	看板・標識など	Scene Text向け

実際のコードが以下です。
下記のHuggingFaceのTrOCRの説明ページのコードを参考に、Google Colabなどで使えるようにNotebook形式で作成しています。 huggingface.co

HuggingFaceではVisionEncoderDecoderModelがTrOCRのモデル呼び出しになります。モデルは看板や標識など現実の画像にある文字に対応したモデルであるmicrosoft/trocr-small-strを使用しました。

import os
os.environ["HF_HOME"] = ".cache"

from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image

# 画像読み込み
image = Image.open("image.png").convert("RGB")

# 画像の表示（Notebookの出力セル内に表示できる）
display(image)

# モデルの準備
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-str')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-str')
pixel_values = processor(images=image, return_tensors="pt").pixel_values

# OCR実行
generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(generated_text)