MMOCRを使ってOCRを作成する（Text Recognation)①

はじめに

コンペティションに参加した時にMMOCRというものを知って、まだ日本語の記事もほとんど作成されていないようだったので、これを使ってOCR作成までの記事を作ってみようと思います。

MMOCRとは

Pytorchベースのオープンソースで公開されているOCR関連のモデルのツールボックスです。
コンフィグファイルを変更することで、各モデルの切り替えを行いながら簡便に色々なモデルを試すことができます。

詳しい説明は、こちらを参照ください。

英数字用の文字列認識用モデルでの学習

今回は、SegOCRというモデルの以下のコンフィグファイルを使って解説を実施していきます。

• seg_r31_1by16_fpnocr_toy_dataset.py

今回は、英数記号のみが認識対象となります。
別途、辞書ファイルを作成することで日本語や漢字、その他記号などを学習させることもできますが、それは次回以降に解説したいと思います。

学習データの準備

まず、学習データとして文字列部分を抜き出した画像ファイル及び、正解の文字列データを準備します。

画像

テキスト
03/09/2009

以下の形式で、ラベル用のテキストファイルを作成します。

label.txt

ファイル名1 テキスト1
ファイル名2 テキスト2
ファイル名3 テキスト3
.
.
.

学習の実施

学習用のコードは、こちらで、公開しています。

上記のコードgoogle collaboratoryで開き、以下の箇所を編集するのみで実施することができます。
（MMOCR側のバージョン指定を行なっていないので、今後のバージョンでは動作しなくなる可能性があります。）

編集する箇所

## アノテーションフォルダの設定
#　元になるテキストファイル

train_label_path = #訓練用ラベルファイルのパス
val_label_path = #検証用ラベルファイルのパス
train_label_lmdb = #訓練用ラベル変換出力したフォルダのパス
val_label_lmdb = #検証用訓練用ラベル変換出力したフォルダのパス

# 画像ファイルのパス
img_path = 画像フォルダのパス
train_img_path = #訓練用画像フォルダのパス
val_img_path = #検証用画像フォルダのパス

# ルートフォルダ
root_path = #ルートフォルダのパス

# 作業フォルダ（学習時のモデル出力先（エポック毎に出力))
work_path = #作業フォルダのパス

この部分のみ変更していただければ、google collaboratory上で学習を実施することができるようにしていますが、上から説明をしていこうと思います。

モジュールのインポート

import torch
import json
from skimage import measure
import numpy as np
print(torch.__version__)
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device = ", device)

MMCV（MMOCRやMMDetectionで使う画像処理関連のツール集）のインストール

#!pip install mmcv-full
!pip install https://download.openmmlab.com/mmcv/dist/cu113/torch1.11.0/mmcv_full-1.4.7-cp37-cp37m-manylinux1_x86_64.whl

今後のGoogle collaboratoryのPytorchやcudaのバージョンが変わった場合は、適切なバージョンに変更する必要があります。

MMDetection（物体検出のツール集）のインストール

# Install mmdetection
!pip install mmdet

MMOCRのインストール

!git clone https://github.com/open-mmlab/mmocr.git
%cd mmocr
!pip install -r requirements/build.txt
!pip install -v -e .  # or "python setup.py develop"

コンフィグファイルの読み込み

%cd /content/mmocr

from mmcv import Config
cfg = Config.fromfile('/content/mmocr/configs/textrecog/sar/sar_r31_parallel_decoder_toy_dataset.py')

パスの設定(解説済み）

## アノテーションフォルダの設定
#　元になるテキストファイル

train_label_path = #訓練用ラベルファイルのパス
val_label_path = #検証用ラベルファイルのパス
train_label_lmdb = #訓練用ラベル変換出力したフォルダのパス
val_label_lmdb = #検証用訓練用ラベル変換出力したフォルダのパス

# 画像ファイルのパス
img_path = 画像フォルダのパス
train_img_path = #訓練用画像フォルダのパス
val_img_path = #検証用画像フォルダのパス

# ルートフォルダ
root_path = #ルートフォルダのパス

# 作業フォルダ（学習時のモデル出力先（エポック毎に出力))
work_path = #作業フォルダのパス

アノテーションファイルの変換

訓練用アノテーションファイルの変換

#訓練ラベルをlmdbに変換
!python /content/mmocr/tools/data/utils/lmdb_converter.py \
    $train_label_path \
    $train_label_lmdb　 --label-only

検証用アノテーションファイルの変換

#検証ラベルをlmdbに変換
!python /content/mmocr/tools/data/utils/lmdb_converter.py \
    $val_label_path \
    $val_label_lmdb　 --label-only

コンフィグの編集

以下のようにコンフィグファイルの設定を変更します。
以下は、初期設定のままの状態で学習、検証用データの部分をオリジナルデータに差し替えたものになります。
用途や環境により、コンフィグファイルの他の項目も変更してみてください。

from mmdet.apis import set_random_seed

# コンフィグを変更する
# segOCR

# 作業フォルダ
cfg.work_dir = work_path

# rootフォルダ
cfg.root = root_path

# アノテーションファイルのパス
cfg.train_anno_file1 = train_label_path
cfg.train_anno_file2 = train_label_lmdb
cfg.img_prefix = train_img_path
cfg.train1.ann_file = train_label_path
cfg.train1.img_prefix = train_img_path
cfg.train2.ann_file = train_label_path
cfg.train2.img_prefix = train_img_path
cfg.test.ann_file = val_label_lmdb
cfg.test.img_prefix = val_img_path
cfg.test_anno_file1 = val_label_lmdb
cfg.train_list[0].ann_file = train_label_path
cfg.train_list[0].img_prefix = train_img_path
cfg.train_list[1].ann_file = train_label_lmdb 
cfg.train_list[1].img_prefix = train_img_path
cfg.test_list[0].ann_file = val_label_lmdb
cfg.test_list[0].img_prefix = val_img_path
cfg.data.test.datasets[0].ann_file = val_label_lmdb
cfg.data.test.datasets[0].img_prefix = val_img_path
cfg.data.train.datasets[0].ann_file = train_label_path
cfg.data.train.datasets[0].img_prefix = train_img_path
cfg.data.train.datasets[1].ann_file = train_label_lmdb
cfg.data.train.datasets[1].img_prefix = train_img_path
cfg.data.val.datasets[0].ann_file = val_label_lmdb
cfg.data.val.datasets[0].img_prefix = val_img_path


# The original learning rate (LR) is set for 8-GPU training.
# We divide it by 8 since we only use one GPU.
cfg.optimizer.lr = 0.001 / 8
cfg.lr_config.warmup = None
# Choose to log training results every 40 images to reduce the size of log file. 
cfg.log_config.interval = 40

# Set seed thus the results are more reproducible
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)


# Let's have a look at the final config used for training
print(f'Config:\n{cfg.pretty_text}')


# 編集したコンフィグをファイル出力
cfg.dump('/content/seg_r31_1by16_fpnocr_toy_dataset.py')

学習の実行

以下のコードで学習を実行します。

from mmocr.datasets import build_dataset
from mmocr.models import build_detector
from mmocr.apis import train_detector
import os.path as osp

# Build dataset
datasets = [build_dataset(cfg.data.train)]

# Build the detector
model = build_detector(
    cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))
# Add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES

# Create work_dir
#mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
train_detector(model, datasets, cfg, distributed=False, validate=True)

予測の実施

以下のコードで予測を実施します。

from mmocr.apis import init_detector, model_inference

img = #確認する画像ファイルのパス
checkpoint = # 利用するモデルのパス
out_file = # 出力するファイルのパス

model = init_detector(cfg, checkpoint, device=device)
if model.cfg.data.test['type'] == 'ConcatDataset':
    model.cfg.data.test.pipeline = model.cfg.data.test['datasets'][0].pipeline

result = model_inference(model, img)
print(f'result: {result}')

img = model.show_result(
        img, result, out_file=out_file, show=False)

以下の形で出力されます。
result:{'text':予測した文字列,'socre':予測結果のスコア}

最後に

本内容は以下のチュートリアルの内容をベースに作成しています。

次回は日本語も含めた文字認識のやり方を解説しようと思います。

OCR関連の投稿記事

MMOCRを使ってOCRを作成する（Text Recognation)②