機械学習OCRモデルのfindtextCenterNetをアップデートしました

bunkoOCRに使用している、機械学習OCRモデルのfindtextCenterNetを新しくしました。
以前のバージョンではTensorflowを用いていたのですが、CoreMLのコンバーターがPyTorchに移行しつつあるのを感じ取ってしまったため、PyTorchに移行することにしました。
学習し直すのに合わせて、モデルの構造もちょっと修正することにしました。

bunkoOCRアプリの方にも更新を反映させますので、もうしばらくお待ちください。

実行結果

入力画像1

コマンド実行の出力

./run_ocr.py testdata/test1.png
---------------------
吾輩は猫である
名前はまだない

手書き文字認識

入力画像2

コマンド実行の出力

./run_ocr.py testdata/test2.png
---------------------
メロスは激怒した。必ず、かの邪智暴虐じゃちぼうぎゃくの王を除かなければならぬと決意した。メロスに
は政治がわからぬ。メロスは、村の牧人である。笛を吹き、羊と遊んで暮して来た。けれど
も邪悪に対しては、人一倍に敏感であった。きょう未明メロスは村を出発し、野を越え山越
え、十里はなれた此このシラクスの市にやって来た。メロスには父も、母も無い。女房も無
い。十六の、内気な妹と二人暮しだ。この妹は、村の或る律気な一牧人を、近々、花婿はなむことし
て迎える事になっていた。結婚式も間近かなのである。メロスは、それゆえ、花嫁の衣裳や
ら祝宴の御馳走やらを買いに、はるばる市にやって来たのだ。先ず、その品々を買い集め、
それから都の大路をぶらぶら歩いた。メロスには竹馬の友があった。セリヌンティウスであ
る。今は此のシラクスの市で、石工をしている。その友を、これから訪ねてみるつもりなの
だ。久しく逢わなかったのだから、訪ねて行くのが楽しみである。歩いているうちにメロス
は、まちの様子を怪しく思った。ひっそりしている。もう既に日も落ちて、まちの暗いのは

モデルの説明

CenterNet https://github.com/xingyizhou/CenterNet という物体認識の手法があり、物体の中心位置である場所の存在確率をヒートマップで出力します。

物体(文字)の中心位置がマークされる

中心位置のマップと、入力画像

中心位置の他に、Boxの縦と横の大きさを、中心位置の場所に出力します。

中心位置に、縦と横のサイズの値が
それぞれ出力される

このCenterNetの手法は、バックボーンのネットワークを選択できます。findtextCenterNetでは、EfficientNetV2 https://github.com/google/automl/tree/master/efficientnetv2 を使用しました。EfficientNetV2-XLを使用しています。

findtextCenterNetの構成図

TextDetectorの部分の構成

findtextCenterNetでは、このCenterNetの手法を用いたTextDetectorの部分と、Transformerの部分の2段階で構成されています。

まずTextDetectorで、画像の中から文字の座標を検出します。その文字Boxを文字列の順に配列します。

文字列のラインを検出する

それぞれの検出できた文字には、100次元の特徴量ベクトルが割り振られます。この特徴量ベクトルは同じ文字では一致するように、異なる文字では異なるように、近いものは近くなるように、学習させます。

それぞれの認識された文字に、特徴量ベクトルが割り振られる

文字列の順番に、この特徴量ベクトルを配列し、Transformerに入力します。

Transformerの部分の構成

Transformerは、青空文庫やWikipediaなどのデータや、スキャンした手持ちの資料を手動で修正したデータ、ランダムに生成した文字列などから、文字コードが正しく生成されるように学習させます。

このモデルの特徴

日本語のOCRでは、縦書きや横書きを認識できる他、ふりがな(ルビ)や圏点を認識して出力することができます。
ふりがなをちゃんと処理できるモデルはなかなか見つからなかったので、作ってみた次第です。