効率的な OCR モデル PGNet
OCR モデル PGNet を解説し、SageMaker でエンドポイントを作ってみます。
とにかく使ってみたい方向けに gist にコードを上げています。
概要
Point Gathering Network (PGNet) は、Point Gathering という演算を使うことで、OCR を実現するための複雑な処理の組み合わせを回避し、高速化を図ったモデルです。AAAI-21 で採択されています。
P. Wang, C. Zhang, F. Qi, S. Liu, X. Zhang, P. Lyu, J. Han, J. Liu, E. Ding, G. Shi
PGNet: Real-time Arbitrarily-Shaped Text Spotting with Point Gathering Network
https://arxiv.org/abs/1507.05717
図1 にその成果が書かれていて分かりやすいですね。SOTAである ABCNet [Liu, et al., CVPR2020] よりも速く動作することが分かります。
従来研究との比較
図2が分かりやすいので書き加えてみました。
まず緑の枠で囲ったネットワークは、物体検出でもよく使われる ROI (Region of Interest) 系の演算を行ってテキスト領域を認識し、その後、実際にテキストを認識するため全体の計算時間が長くなります。論文中では two-stage の手法と呼んでいます。
一方、提案手法はテキスト領域の認識とテキストの認識を一括で行う方法を提案しています。これはオレンジで囲んだ CharNet [Xing et al., ICCV2019] と同じ枠組みです。しかし、CharNet は図中の GT: W, C のところに着目すると、文字単位 (C) でのアノテーションが必要であり、一方 PGNet は単語単位 (W) でのアノテーションのみで良いことからデータ準備の面で有利です。
PGNet の全体像
図3から説明していきましょう。まず図左にあるように画像から特徴抽出を行います。特徴抽出には物体検出によく使われるFPN（Feature Pyramid Networks）が利用されます。抽出した特徴 $F_{visual}$ は元の画像サイズの 1/4 のサイズになります。続いてこの $F_{visual}$ から以下の情報を求めるネットワークを学習します。図で上から順だと説明しにくいので、少し順番をかえて説明します。
- Text Center Line (TCL)
TCL はテキストの中心領域を表していて、テキストの中心領域かどうかの1次元スコアが各ピクセルに入っています（1チャネル） - Text Border Offset (TBO)
TCL の各ピクセルからテキスト領域の上側の点までの距離 (x方向とy方向)と下側の点までの距離(x方向とy方向) が、各ピクセルに入っています (4チャネル) - Text Direction Offset (TDO)
テキストの方向が入っています。2次元で方向を表すので2チャネルです。 - Text Character Classification Map (TCC)
各ピクセルがどの文字を表しているかを示します。37文字の英数字なら37チャネルです。
これらを推定するネットワークは非常に簡単で、それぞれConv+Batch Normalization になっています。詳細は github のコード を見てみてください。
このうち TBO はそのままテキスト領域の検出に使われ、それ以外の部分はテキスト認識に利用されます。テキスト認識は次で説明する Point Gathering で行います。
Point Gathering (PG)
あるテキスト領域に対する PG について、論文中では以下のような式で説明されています。
$$
P_\pi = gather(TCC, \pi)
$$
$\pi$ は TCL で予測した N 個の点を TDO の方向に従って並べ替えたもの、TCC は上で求めた各文字の確率に関するマップです。これらを使って、N 個の点に対応する各文字の確率 $P_\pi$ を求めます。gather というのがよくわかりませんでしたが、 こちらのコードを参考にすると、N 個の点の確率分布を TCC から取ってくるだけのようです。TCC はConv+Batch Normalizationの結果をそのまま使う方法だけでなく、$F_{visual}$ と同時に Graph Convolution を使って精度を改善する方法も提案されています (3.5節)。
学習時にはこの$P_\pi$と正解のテキスト $L_i$ との比較でロスを計算します。M個のテキストが認識されたら、それぞれのロスの総和をとります。
L_{PG−CTC}=∑_{i=1}^M CTC\_loss(P_{\pi_i}, L_i)
ここでは CTC (Connectionist Temporal Classification) ロスというのを使っています。CTC ロスの説明は towardsdatascience.com がわかりやすいと思います。なぜ CTC を使うかというと、実は学習データが文字単位でのアノテーションを行っていないので、文字を比較してロスを計算することができないからです。例えば、OCR という単語を認識しようとしたとき、"OCC R"みたいに認識されるかもしれません。CC のように同じ文字が連続して認識されたり、空白ができたりしてしまいます。CTC ロスは、重複する単語を除去したり、空白を除去したりする CTC 特有の操作 (CTC Decoding) を想定したロスになっています。この $L_{PG−CTC}$ と、それ以外に TCL, TBO, TDO のためのロスも計算して、それぞれを重み付けしたロスでネットワーク全体を学習しています。
推論するときは gather 関数を使って $P_\pi$ を求めたら、それに対して、重複する単語を除去したり、空白を除去したりする CTC Decoding をかけるだけです。
R_\pi = CTC\_Decoder(P_\pi)
SageMaker へデプロイして日本語OCRを作る
PGNet を開発している Baidu が日本語を含む多言語OCRモデルを公開しており、PaddleOCRという形で簡単に利用できるようにしています。
PaddleOCR
https://github.com/PaddlePaddle/PaddleOCR
実はフレームワークは TensorFlow とかではなく、Baidu 独自の Paddle (PArallel Distributed Deep LEarning) とよばれるフレームワークです。とはいえコンテナイメージまで公開されており、python での API も用意されているので SageMaker に乗せることはできそうです。
以下では SageMaker ノートブックインスタンスで実行することを想定しています。インスタンスタイプは t2.medium で良いですが、コンテナイメージのビルドに容量を使うので、ディスクサイズを 50GB くらいにしておいてください。
例によって面倒な方は gist のノートブックをダウンロードしてそのまま実行できます。
PaddleOCR を動かす SageMaker コンテナイメージの準備
ビルドを home で行うように設定
これからビルドする Docker イメージはなかなか大きく、デフォルトの /var/lib/docker ではディスクサイズが足りずビルドに失敗します。そこで以下を最初に実行して、home でビルドするようにしましょう。
!sudo /etc/init.d/docker stop
!sudo mv /var/lib/docker /home/ec2-user/SageMaker/docker
!sudo ln -s /home/ec2-user/SageMaker/docker /var/lib/docker
!sudo /etc/init.d/docker start
コンテナイメージ作成に必要なファイル群
今回は Paddle というフレームワークを使うので、SageMaker で用意したコンテナではなく自作して使います。作り方は以下のサンプルが参考になります。
必要なファイルは以下のとおりです。
- Dockerfile
- GPU 上で Paddle が使えるよう Cuda を利用できる環境を用意
- Python からも利用できるよう pip で Paddle の Python ライブラリをインストール
- それ以外は上記のサンプルに従って、必要なデプロイ用のライブラリをインストール
- dockerd-entrypoint.py
- 特に編集不要
- model_handler.py を実行するように実装されている
- model_handler.py
- 日本語の OCR モデルをロードする処理を実装する
- 画像を受け取ったら、ロード済みのモデルで推論して、結果を返す実装をする
Dockerfile
長くなるので全体は gist を見てください。
今回は GPU を使って推論するイメージを作成します。PaddleOCR を GPU で動かすためには以下が必要です。
- Cuda 10.2
- cuDNN v7.6
- glibc 2.2.3 (Ubuntu 16.04 で利用可能)
そこで、これを利用可能なベースイメージを取得するよう変更しましょう。ubuntu:16.04のままだと CPU しか使うことができません。
FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu16.04
Python のバージョンを上げておきます。
ARG PYTHON_VERSION=3.7.10
あとは python から利用できるよう paddle のライブラリを pip でインストールします。
# install paddleocr
RUN pip3 install paddlepaddle-gpu==2.0.2
RUN pip3 install "paddleocr>=2.0.1"
dockerd-entrypoint.py
こちらは変更点がないので説明を省略します。
model_handler.py
実際にモデルで推論する際の挙動を記述します。その挙動は、グローバルの関数の handle(data, context) で書かれているように
- 初期化されていなければ initialize で初期化する (モデルをロードするなど）
- データがくればhandleで予測した結果を返す
となっています。つまり、initialize と handle を実装すればOKです。initialize は以下のようにモデルをロードする処理を書きます。SageMaker は S3 にモデルをおいておけば、それを /opt/ml/model/ においてくれますので、そこに置かれた各種モデルを読むようにします。S3へのモデル配置はこのあと行います。
self.initialized = True
model_dir = "/opt/ml/model/"
# Load ocr model
try:
self.ocr = PaddleOCR(det_model_dir=os.path.join(model_dir,'model/det'),
rec_model_dir=os.path.join(model_dir,'model/rec/ja'),
rec_char_dict_path=os.path.join(model_dir,'model/dict/japan_dict.txt'),
cls_model_dir=os.path.join(model_dir,'model/cls'),
use_angle_cls=True, lang="japan", use_gpu=True)
handle は preprocess と inference にかけます。
def handle(self, data, context):
model_input = self.preprocess(data)
model_out = self.inference(model_input)
return model_out
それぞれの実装は以下のとおりです。preprocess では base64 形式のデータを想定して、1枚の画像を numpy array に変換してリストにしています。これは PaddleOCR が numpy を受け取るのでそのようにしています。ocr.ocr で実行した予測結果 result は、[テキスト領域の座標群, (テキスト、スコア）] となります。このうち、スコアはfloat32で json に変換できないので float64に直してリスト形式にしておきます。
def preprocess(self, request):
# Take the input data and pre-process it make it inference ready
img_list = []
for idx, data in enumerate(request):
# Read the bytearray of the image from the input
img_arr = data.get('body')
img_arr = base64.b64decode(img_arr)
img_arr = Image.open(BytesIO(img_arr))
img_arr = np.array(img_arr)
# Check the number of dimension
assert len(img_arr.shape) == 3, "Dimension must be 3, but {}".format(len(img_arr.shape))
img_list.append(img_arr)
return img_list
def inference(self, model_input):
res_list = []
# Do some inference call to engine here and return output
for img in model_input:
result = self.ocr.ocr(img, cls=True)
for res in result:
## because float32 is not json serializable, score is converted to float (float64)
## However the score is in tuple and cannot be replaced. The entire tupple is replaced as list.
string = res[1][0]
score = float(res[1][1])
res[1] = [string,score]
res_list.append(result)
return res_list
これを一発で書けるか？といわれると、そんな人は多くいないと思います。デプロイして画像を送ってエラーがないかを確認して、エラーが出れば上記を修正する、という作業を繰り返す必要が出てきます。ライブラリが必要であれば Dockerfile に追記、推論のスクリプトにバグがあればmodel_handler.py を修正します。初回のコンテナイメージのビルドには数十分の時間がかかりますが、一度ビルドするとキャッシュされるので、細々とした変更に多くの時間はかかりません。
OCR モデルの S3 へのアップロード
PaddleOCR のライブラリは、モデルがローカルにない場合は自動でモデルをダウンロードする仕組みをもっています。しかし今回は、事前にモデルをダウンロードして S3 に保存しておき、デプロイするときは S3 からモデルをダウンロードして利用するようにします。自分で管理するので手間ではありますが、もしライブラリがモデルのダウンロードに失敗しても S3 からダウンロードして対応できます。
以下のテキストエリア検出、テキスト認識（認識モデルと日本語辞書）、テキスト分類（角度調整等に利用）の4つのファイルがあり、それぞれダウンロード・解凍して、以下に示すフォルダ構造で保存します。このフォルダ構造は、コンテナイメージをビルドする際に作成した model_handler.py の関数 initialize でモデルを読み込む際の階層構造と合わせる必要があります。また、モデルの URL はこちらのファイルから確認できます。テキスト認識にかかる2つのファイルは言語固有で、それ以外は、各言語共通のものを使います。
以下の階層構造でモデルなどを保存したら、それらを1つに圧縮して model.tar.gz にします。これは SageMaker が tar.gz で圧縮されていることを前提とするためです。エンドポイントを作成する際は、このmodel.tar.gzは解凍されて /opt/ml/model に展開されます（つまり /opt/ml/model/model/det/ 以下にテキストエリア検出モデルが展開されます）。
model
├── det ... (テキストエリア検出)
├── rec
│ └── ja ... (テキスト認識モデル)
├── dict ... (日本語辞書)
└── cls ... (テキスト分類)
エンドポイントの作成
SageMaker の Model を作成して deploy を実行すればエンドポイントを作成できます。image_uriには ECR に push したコンテナイメージの URI を、model_data には先ほどアップロードした OCR モデル (model.tar.gz) へのパスを渡します。
from sagemaker.model import Model
from sagemaker.predictor import Predictor
ocr_model = sagemaker.model.Model(image_uri,
model_data=model_uri,
predictor_cls=Predictor,
role=sagemaker.get_execution_role())
今回は GPU を使いたいので ml.g4dn.xlarge をインスタンスに選んでデプロイします。ただし、これには時間がかかるので、もしデバッグがすんでいなければ、instance_type="local"と指定して、ローカル環境でまずはテストしてみましょう。以下でエラーが出る場合は、コンテナイメージに不具合があるはずなので、エラーを CloudWatch Logs で確認して、Dockerfile や model_handler.py を修正してコンテナイメージをもう一度ビルドします。
predictor = ocr_model.deploy(initial_instance_count=1,instance_type="ml.g4dn.xlarge")
画像を OCR にかけてみる
某所から持ってきた画像を試してみます。画像は model_handler.py で実装したように base64 でエンコードして送ります。
%%time
import base64
from io import BytesIO
from PIL import Image
import PIL
image = Image.open("test01.png")
if isinstance(image,PIL.PngImagePlugin.PngImageFile):
image = image.convert('RGB')
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue())
import json
response = json.loads(predictor.predict(img_str))
response には結果が入っていますが、わかりやすくするため、検出された場所に赤枠で囲み、認識されたテキストと対応付けるための通し番号を付けます。
import numpy as np
import cv2
x_offset = 20
y_offset = 0
for i, res in enumerate(response):
box = np.reshape(np.array(res[0]), [-1, 1, 2]).astype(np.int64)
image = cv2.putText(np.array(image), '('+str(i)+')', (box[0][0][0] -x_offset, box[0][0][1]-y_offset),
cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), 1, cv2.LINE_AA)
image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
最後に表示しましょう。
import matplotlib.pyplot as plt
for i,res in enumerate(response):
print('('+str(i)+'): '+res[1][0], end=', ')
fig = plt.figure(dpi=200)
plt.imshow(image)
plt.show()
結果はこのような感じです。改行のある文章が各行ごとに検出されていますが、全体の日本語の精度はなかなか良さそうに見えます。一方、日本語にしているせいなのか、どうも英数字の検出精度がいまいちに見えます (Amazon EC2 -> AmaJnＥC2 など)。また、数字は小数点が検出されていません。
また、検出速度ですが、前者の画像 (603 × 434 ピクセル）は概ね 130 msec での検出、後者の画像 (910 × 558 ピクセル) は 260 msec での検出となりました。
さいごに
このノートブックを試した方はエンドポイントを削除するのを忘れないようにしましょう。