Microsoftが提供するTransformerベースのOCRであるTrOCRの導入手順を解説します。

TrOCRとは

TrOCRとはMicrosoftが発表したTransformerベースのOCRです。[参考:arxiv]

従来のAIOCRは、画像中の文字を検出する文字検出にCNNを、文字認識にRNNを適用してモデルを構築することが一般的でした。
一方でTrOCRでは文字検出と文字認識共に、Transformerを適用しています

Transformerを適用することによる利点は下記3点であると論文で語られています。

1. 事前にトレーニングされたImageTransformerモデルと、TextTransformerモデルを使用し、精度を確保
2. バックボーンに畳み込みネットワークを必要としないためモデルの実装と保守の簡素化
3. OCRのベンチマークデータセット(IAM(手書き)/SROIE(活字))にて最先端の精度

簡単にまとめると、Transformerの適用により、事前学習済みモデルによって精度を確保し、モデルが簡素になっているOCRであると言えそうです。

一方でデメリットとしては、モデルのサイズが大きくなりがちといったところでしょうか？
提供されているTrOCR-Large-IAMは6.4GBにも及びます。

TrOCRの概念図 出典:https://arxiv.org/pdf/2109.10282.pdf

TrOCRの導入手順

TrOCRのセットアップ手順

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

それでは、TrOCRをインストールしていきます。他の機械学習環境に影響を与えないためにMinicondaの仮想環境上に構築していきます。Minicondaのインストール手順は公式ドキュメントをご参照ください。

TrOCRはmicrosoftが提供するUniLM AIというプロジェクトの一つです。
UniLM AIは要約抽出や、OCR、翻訳などの様々な事前トレーニング済みモデルを提供しているGitのプロジェクトです。

# python version 3.7のconda環境を作成
$ conda create -n trocr python=3.7

# 作成した環境をアクティベート
$ conda activate trocr

$ cd trocr

# gitからコードをclone
$ git clone https://github.com/microsoft/unilm.git

# trocrのディレクトリに移動
$ cd unilm/trocr/

# Proxy配下の場合は事前にproxyを設定
$ export http_proxy="http://"username":"password"@proxy:port"
$ export https_proxy="http://"username":"password"@proxy:port"

# 必要なライブラリをインストール
$ pip3 install pybind11
$ pip3 install -r requirements.txt
$ pip3 install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" 'git+https://github.com/NVIDIA/apex.git'

以上で、TrOCRのセットアップは完了です。
続いてTrOCRの学習済みモデルを使って文字認識していきます。

TrOCRで文字認識

学習済みモデルが公開されているため、このモデルを使って手書き英語の文字認識していきます。
なお、学習済みモデルが非常に大きいためダウンロードに時間がかかります。(およそ20分)

# モデルダウンロード用ディレクトリ作成
$ mkdir models
$ cd models
# 学習済みモデルダウンロード
$ wget https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt
$ cd ..

# テスト用画像準備
$ mkdir test_img
# test_img直下に手書き英語を含む画像を配置

今回は以下のような手書き英語の公開データセットを用意しました。

それでは、文字認識を実行していきます。

#学習済みモデルをダウンロードします
$ wget https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt

trocr/pic_inference.pyにtest用画像とモデルのパスを設定します。

if __name__ == '__main__':
    model_path = 'path/to/model'
    jpg_path = "path/to/pic"
    #以下、設定例
    #model_path = "home/miniconda3/envs/trocr/unilm/trocr/models/trocr-base-handwritten.pt"
    #jpg_path = "home/miniconda3/envs/trocr/unilm/trocr/test_img/test_image.png"

trocr/pic_inference.pyをpython3 pic_inference.pyで実行します

$ python3 pic_inference.py

| [label] load dictionary: 50265 types
Marquette . # ←文字認識出力結果
done

いかがですか？概ね認識されているかと思います。

urllib.error.HTTPError: HTTP Error 404: Not Found

もし、以下のような404エラーが出力される場合は、ベースとなるモデルのロードに失敗しています。

$ python3 pic_inference.py
Downloading: "https://github.com/pytorch/fairseq/archive/master.zip" to ...中略
raise HTTPError(req.full_url, code, msg, hdrs, fp)
urllib.error.HTTPError: HTTP Error 404: Not Found

下記の通り、ローカルにモデルをダウンロードしておき、ファイルを修正することで解消します。

# モデルのダウンロード
$ wget https://dl.fbaipublicfiles.com/fairseq/models/roberta.large.tar.gz
$ tar -zxvf roberta.large.tar.gz

#trocr/deit_models.pyの以下を編集
elif args.decoder_layers == 12:
	logger.info('Load pre-trained decoder parameters from roberta.large')
	# 下記のモデルが404のためローカルにダウンロードした後にロード
	# roberta = torch.hub.load('pytorch/fairseq', 'roberta.large')
	from fairseq.models.roberta import RobertaModel
	roberta = RobertaModel.from_pretrained(
		'roberta.large',
		'model.pt',
		'/home/miniconda3/envs/trocr/unilm/trocr/models/roberta.large')

まとめ

本記事ではTransformerベースのOCRであるTrOCRを紹介させて頂きました。
ソースコードをクローンしていただくと分かりますが、コードの量が非常に少ないことが分かります。一方で、モデルサイズが1GBを超え非常に巨大であることも分かります。

CNN, RNNなどを必要としないためコードが簡素であるため見通しが良く、
どこでどのような処理を行っているか確認しやすい点が良いですね。
一方で、モデルが非常に大きいため、組み込み機器等、ハードウェア制約がある環境では、動作困難な可能性があります。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1. TrOCR: Transformer-based Optical Character Recognition with Pre-trained Models

2. microsoft/unilm