AutoGluonを使ってMMOCRの推論が行えるよ

AutoGluon OCR PyTorch

公開日：2022年9月27日
最終更新日：2022年11月17日


この記事は2022年11月17日AutoGluon 0.6が公開されたのに合わせて更新されています。

はじめに

以前「MMOCR」の推論を行う記事を書きました。
touch-sp.hatenablog.com
今回は「AutoGluon」から「MMOCR」を使用する方法を紹介します。

Pythonスクリプト

非常にシンプルなスクリプトです。

サンプル画像のダウンロードも自動的に行われます。
学習済みモデルの事前準備（ダウンロード）も必要ありません。

from autogluon.multimodal import download, MultiModalPredictor
import mmocr.datasets.pipelines

url = "https://raw.githubusercontent.com/open-mmlab/mmocr/main/demo/demo_text_recog.jpg"
mmocr_image_name = download(url)

checkpoint_name = 'abinet_academic'

predictor = MultiModalPredictor(
        hyperparameters={
            "model.mmocr_text_recognition.checkpoint_name": checkpoint_name,
        },
        problem_type="ocr_text_recognition",
    )

pred = predictor.predict({"image": [mmocr_image_name]})

print(pred)

結果

[['star'], [0.999382957816124]]

その他のモデル

「abinet_academic」以外に以下のモデルが使用できるようです。

"abinet_academic"
"sar_r31_parallel_decoder_academic"
"seg_r31_1by16_fpnocr_academic"
"nrtr_r31_1by16_1by8_academic"

環境構築

Ubuntu 20.04 on WSL2
CUDA 11.6.2
python 3.8.10



Pythonのライブラリはすべてpipで可能です。

pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 torchtext==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu116
pip install mmcv-full==1.7.0 -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.12.0/index.html
pip install mmdet==2.25.3
pip install mmocr==0.6.3
pip install autogluon==0.6.0