NERとは

Named Entity Recognitionの頭文字を取ったもので、文中から固有表現 (Named Entity) を抽出し、それを固有名詞（人名、組織名、地名など）や日付、時間表現、数量、金額、パーセンテージなどのあらかじめ定義された固有表現分類へと分類するもの[wiki参照]を指します

NERを試してみる

MMOCRで配布されているNERのdemoを動かしてみます。
そのために少し準備が必要です。
MMOCRのサイトからNERの学習済みcheckpoint、Datasetをダウンロードし、これを以下のように配置します。

mmocr
    ├-- mmocr
    ├-- demo
        ├-- ner_demo.py
    ├-- configs
        ├-- ner
            ├-- bert_softmax
                ├-- bert_softmax_cluener_18e.py
    ├-- checkpoints(なければmkdir)
        ├-- bert_softmax_cluener-eea70ea2.pth(これを配置)
    ├-- data
        ├-- cluener2020(これを配置)
            ├-- cluener_predict.json
            ├-- dev.json
            ├-- README.md
            ├-- test.json
            ├-- train.json
            ├-- vocab.txt

次に以下のコマンドを実行します。

python demo/ner_demo.py \
    configs/ner/bert_softmax/bert_softmax_cluener_18e.py \
    checkpoints/bert_softmax_cluener-eea70ea2.pth

実行すると入力を求められるのでNERを試したい文章を入力します。
ただしこのモデルはCLUENER2020という中国語のデータセットで学習させたものですので中国語（CLUENER2020のサイトに載っていた文）を入力してみます。

Please enter a sentence you want to test: 北京勘察设计协会副会长兼秘书长周荫如
organization: 北京勘察设计协会
position: 副会长
position: 秘书长
name: 周荫如

組織名や役職、人物名に分類されています。
ここで少し述べておくと、CLUENER2020は以下のラベルに分類するデータセットになっています。

address, book, company, game, government, movie, name, organization, position, scene

Fine-Tuning

MMOCRはNERのためのモデルとしてBERTのみを提供しています。
以下では付属のDataset、任意のDatasetでこの既存モデルをFine-Tuningする方法について記します。

MMOCR付属のDatasetを用いたFine-Tuning

MMOCRが用意しているDatasetはCLUENER2020のみです。このDatasetを用いたFine-Tuningについては既にconfig fileが用意されていますので、それを使ってFine-Tuningしてみます。

python tools/train.py \
    configs/ner/bert_softmax/bert_softmax_cluener_18e.py \
    --cfg-options work_dir="work_dirs/cluener"

上記のコマンドを実行することで学習を開始することが出来ます。work_dirに設定したファイルに学習結果の重み、logなど学習結果が格納されます。

python demo/ner_demo.py \
    work_dirs/cluener/bert_softmax_cluener_18e.py \
    work_dirs/cluener/latest.pth

で学習結果を確認できます。
（tools/test.pyの動作は確認できていないです）

任意のDatasetを用いたFine-Tuning

Datasetの準備

任意のDatasetを用いる場合は以下のようなファイルが必要になります。

train.json
dev.json
test.json(使われていない？)
vocab.txt

CLUENER2020のDatasetを元にそれぞれのファイルの構造を見てみます。

train.json

...
{"text": "生生不息CSOL生化狂潮让你填弹狂扫", "label": {"game": {"CSOL": [[4, 7]]}}}
{"text": "那不勒斯vs锡耶纳以及桑普vs热那亚之上呢？", "label": {"organization": {"那不勒斯": [[0, 3]], "锡耶纳": [[6, 8]], "桑普": [[11, 12]], "热那亚": [[15, 17]]}}}
...

dev.json

...
{"text": "光大三家银行信用卡专业人士分别就境外购买1000欧元、1000英镑和1万港币的商品，", "label": {"company": {"光大": [[0, 1]]}}}
{"text": "4、雷吉纳vsac米兰推荐：0", "label": {"organization": {"雷吉纳": [[2, 4]], "ac米兰": [[7, 10]]}}}
...

test.json

...
{"id": 6, "text": "央视新址文化中心外立面受损严重"}
{"id": 7, "text": "单看这张彩票，税前总奖金为5063992元。本张票面缩水后阿森纳的结果全部为0，斯图加特全部为1，"}
...

train.json、dev.jsonには文とそれに対するラベルが対応するアノテーションファイルとなっています。構造としては辞書の入れ子となっていることがわかります。一方test.jsonには文が羅列されています。
（ぱっと見ですがtestとdevの命名が逆な気がします）

vocab.txt

...
request
##gence
qt
##っ
1886
347
363
q7
##zzi
...

また、vocab.txtは語彙の羅列になっています。

wiki Dataset

これらの構造を考えながら日本語のデータセットであるner-wiki-datasetで学習を行ってみます。
ner-wiki-datasetからner.jsonをダウンロードしてきて、その構造を確認します。

ner.json

[
    {
        "curid": "3572156",
        "text": "SPRiNGSと最も仲の良いライバルグループ。",
        "entities": [
            {
                "name": "SPRiNGS",
                "span": [
                    0,
                    7
                ],
                "type": "その他の組織名"
            }
        ]
    },
    {
        "curid": "2415078",
        "text": "レッドフォックス株式会社は、東京都千代田区に本社を置くITサービス企業である。",
        "entities": [
            {
                "name": "レッドフォックス株式会社",
                "span": [
                    0,
                    12
                ],
                "type": "法人名"
            },
            {
                "name": "東京都千代田区",
                "span": [
                    14,
                    21
                ],
                "type": "地名"
            }
        ]
    },
    ...
]

GithubのREADMEによると以下のような構造をしているらしいです。

curidはデータ元のWikipediaのページID
textはタグ付を行う対象のテキスト
entitiesは固有表現のリスト
nameは固有表現名
spanはtextでの位置
typeは固有表現のタイプ

train.jsonと比較してみると

Wiki Dataset	Cluener2020(MMOCR)
text	text
enitities	label
{"name": , "span": , "type", }	{"type": {"name": ["span"]}}

となっています。またtrain.jsonはJSON形式ですがner.jsonはNDJSON形式です。
これらの違いを踏まえて Wiki DatasetをMMOCRで利用できるように変換します。

converter.py

import random
import json
import ndjson
import sys

def convert_from_wiki_dataset(wiki_dataset_file, train_ratio=0.9):
    '''
    construct dataset from https://github.com/stockmarkteam/ner-wikipedia-dataset
    Parameters:
    -----------
    wiki_dataset_file: str
        path to the json file
    train_ratio: float
        separate dataset into train and test with this ratio
    -----------
    '''
    with open(wiki_dataset_file, 'r', encoding='UTF-8') as f:
        annotated_texts = json.load(f)
    results = []
    for annotated_text in annotated_texts:
        result = {'text': None, 'label': None}
        text = annotated_text['text']
        result['text'] = text

        entities = annotated_text['entities']
        label = {}
        for entity in entities:
            entity_type = entity['type'] # str ex. 施設名, イベント名
            name = entity['name'] # str ex. 日本橋支店, ウッドストック・フェスティバル
            span = entity['span'] # list ex. [49, 54], [24, 39]
            span[1] -= 1
            tmp_dict = {name: [span]}
            if entity_type in label:
                label[entity_type].update(tmp_dict)
            else:
                label[entity_type] = tmp_dict
        result['label'] = label
        results.append(result)
    
    random.shuffle(results)
    train_data = results[0:int(len(annotated_texts)*train_ratio)]
    test_data = results[int(len(annotated_texts)*train_ratio):]

    with open('train.json', 'w', encoding='UTF-8') as f:
        ndjson.dump(train_data, f, ensure_ascii=False)
    with open('test.json', 'w', encoding='UTF-8') as f:
        ndjson.dump(test_data, f, ensure_ascii=False)

def main():
    convert_from_wiki_dataset(sys.argv[1])

if __name__ == '__main__':
    main()

上記のコードを実行すると

python converter.py <ner.jsonのパス>

9:1に分割されたtrain.json、test.jsonが出力されます。
出力されたtrain.jsonを見てみるとMMOCRで利用できる形に変換されていることがわかります。

train.json(wiki dataset)

{"text": "標高1,600-2,200メートルの山地にある森林に生息する。", "label": {}}
{"text": "2017年から第5代メタルワン代表取締役社長を務め、2018年には住友商事との間で国内鋼管事業統合により新会社住商メタルワン鋼管の設立を行うことなどで合意した。", "label": {"法人名": {"メタルワン": [[10, 14]], "住友商事": [[33, 36]], "住商メタルワン鋼管": [[55, 63]]}}}
...

これらのファイルを以下のように配置します。

mmocr
    ├-- mmocr
    ├-- demo
    ├-- configs
        ├-- ner
            ├-- bert_softmax
                ├-- bert_softmax_wiki_18e.py(後ほど配置)
    ├-- checkpoints
    ├-- data
        ├-- wiki(これを配置)
            ├-- test.json
            ├-- train.json
            ├-- vocab.txt

Cluener2020のvocab.txtには日本語の語彙も含まれているため、同じものを利用します。

config fileの記述

次にconfig fileを記述していきます。
使うモデルがCLUENER2020のときと同じことからconfigs/ner/bert_softmax/bert_softmax_cluener_18e.pyを継承してconfigs/ner/bert_softmax/bert_softmax_wiki_18e.pyを作成します。

bert_softmax_wiki_18e.py

_base_ = [
    './bert_softmax_cluener_18e.py'
]

categories = [
    '人名', '法人名', '政治的組織名', 'その他の組織名', '地名', '施設名', '製品名', 'イベント名'
]

test_ann_file = 'data/hogehoge/test.json'
train_ann_file = 'data/hogehoge/train.json'
vocab_file = 'data/hogehoge/vocab.txt'

max_len = 256

上記の`categoriesにはWiki Datasetのラベルをすべて列挙します。

タイプ	固有表現数	備考
人名	2980
法人名	2485	法人又は法人に類する組織
政治的組織名	1180	政治的組織名、政党名、政府組織名、行政組織名、軍隊名、国際組織名
その他の組織名	1051	競技組織名、公演組織名、その他
地名	2157
施設名	1108
製品名	1215	商品名、番組名、映画名、書籍名、歌名、ブランド名等
イベント名	1009

test_ann_file、train_ann_fileには変換によって得られたtest.json、train.jsonを設定します。
max_lenにはtextの入力として何文字までを許容するかを記します。今回であれば大凡400文字がtrain.json内の最大文字数となります。ただし、ほとんどのtextがそれほど大きくないため今回はmax_lenを256に設定してあります。

実行時に256以上のtextが入力されるとWarningを吐きますが、この入力はスキップされ処理は問題なく継続されます。

train.pyの実行

python tools/train.py \
    configs/ner/bert_softmax/bert_wiki_18e.py \
    --cfg-options work_dir="work_dirs/wiki"

上記のコマンドを実行することで学習を開始することが出来ます。

学習結果の確認

python demo/ner_demo.py \
    work_dirs/hogehoge/bert_softmax_wiki_18e.py \
    work_dirs/wiki/latest.pth

で学習結果を確認できます。
試しに適当な分を入力してみるとしっかりNERされています。

Please enter a sentence you want to test: イギリスはリーマンショック直後の2008年10月にイングランド銀行のバランスシートを一気に3倍近く増やした後、2008年11月から2009年3月にかけて段階的に縮小させていった。
地名: イギリス
イベント名: リーマンショック
政治的組織名: イングランド銀行

番外編

NERとOCRの接続

MMOCRのNERはMMOCRのDET、RECOGとは接続されていません。（NERはOCRに直接関係ないからでしょうか）
これをうまく接続することで画像の入力に対し固有表現抽出できるようになります。
NER、DET、RECOGを実行するためのコードは以下のファイルに記述されています。
・NER : demo/ner_demo.py
・DET、RECOG : mmocr/utils/ocr.py

コードの解析

demo/ner_demo.pyはconfigとcheckpointをコマンドライン引数として取り、その後入力された文に対して固有表現抽出を行うものとなっており、簡単なコードです。
一方、mmocr/utils/ocr.pyは複雑な構造をしているため動作を簡単に解説します。
mmocr/utils/ocr.pyの主要なコマンドライン引数は以下の通りになっており、この引数の有無で動作が変化します。

引数	説明
--det	MMOCRが用意しているdetectionモデル名
--det-config	MMOCRが用意しているdetectionモデルについて独自で学習させたconfigファイルのパス
--det-ckpt	MMOCRが用意しているdetectionモデルについて独自で学習させたckptファイルのパス
--recog	MMOCRが用意しているrecogモデル名
--recog-config	MMOCRが用意しているrecogモデルについて独自で学習させたconfigファイルのパス
--recog-ckpt	MMOCRが用意しているrecogモデルについて独自で学習させたckptファイルのパス
--kie	MMOCRが用意しているkieモデル名
--kie-config	MMOCRが用意しているkieモデルについて独自で学習させたconfigファイルのパス
--kie-ckpt	MMOCRが用意しているkieモデルについて独自で学習させたckptファイルのパス

--**-config、--**-ckptは独自で学習させた結果を用いる際のオプションとなります。
また、MMOCRが用意していないモデルの使用は想定されていません。
これらの引数の有無とその動作、呼び出される関数を以下に示します。

引数	動作
det、recogどちらかのみ指定	指定されたどちらかの操作のみを実行、single_inference()
det、recogを指定	det後にrecogを実行、det_recog_kie_inference()
det、recog、kieを指定	det、recog、kieを順に実行、det_recog_kie_inference()

これらのことから以下を実装すればNERとOCRを接続できます。

コマンドライン引数に--ner、--ner-config、--ner-ckptを追加（独自モデルの利用を想定（--nerがNoneでなければOCR結果をNERする））
コマンドライン引数の処理を--ner、--ner-config、--ner-ckptについて追加
demo/ner_demo.pyとdet_recog_kie_inference()をもとにdet_recog_ner_inference()を実装
コマンドライン引数の有無による処理分岐を実装

実装

上記の4項目を実装します。
まずは、1.を実装します。引数を指定しているparse_args()に以下のように追記します

ocr.py

def parse_args():
    ...
    parser.add_argument(
        '--ner',
        type=str,
        default='',
        help='Pretrained named entity recognition algorithm')
    parser.add_argument(
        '--ner-config',
        type=str,
        default='',
        help='Path to the custom config file of the selected ner model. It'
        'overrides the settings in ner')
    parser.add_argument(
        '--ner-ckpt',
        type=str,
        default='',
        help='Path to the custom checkpoint file of the selected ner model. '
        'It overrides the settings in ner')
    ...

次に2.を実装します。コマンドラインの処理は__init__()に追記します。

ocr.py

def __init__(self,
             ...
             kie_config='',
             kie_ckpt='',
             ner='', # 追記
             ner_config='', # 追記
             ner_ckpt='', # 追記
             config_dir=os.path.join(str(Path.cwd()), 'configs/'),
             device=None,
             **kwargs):
    ...
    self.td = det
    self.tr = recog
    self.kie = kie
    self.ner = ner # 追記
    self.device = device
    ...
    self.kie_model = None
    if self.kie:
        ...
    
    self.ner_model = None # 追記
    if self.ner: # 追記
        self.ner_cfg = Config.fromfile(ner_config) # 追記
        # build the model from a config file and a checkpoint file
        self.ner_model = init_detector(ner_config, ner_ckpt, device=self.device) # 追記

次に3.を実装します。
det_recog_kie_inference()をコピーし、以下のように変更を加えます。

ocr.py

def det_recog_ner_inference(self, det_model, recog_model, ner_model): # 変更
    ...
    for filename, arr, bboxes, out_file in zip(self.args.filenames,
                                               self.args.arrays,
                                               bboxes_list,
                                               self.args.output):
        ...
        if self.args.merge:
            img_e2e_res['result'] = stitch_boxes_into_lines(
                img_e2e_res['result'], self.args.merge_xdist, 0.5)
        
        # 以下を追記
        for i in range(len(img_e2e_res['result'])):
            text = img_e2e_res['result'][i]['text']
            ner_result = text_model_inference(ner_model, text)
            print('')
            print('recognized text in box:', text)
            # show the results
            for pred_entities in ner_result:
                for entity in pred_entities:
                    print(f'{entity[0]}: {text[entity[1]:entity[2] + 1]}')
                    if entity[0] in img_e2e_res['result'][i]:
                        img_e2e_res['result'][i][entity[0]].append(text[entity[1]:entity[2]])
                    else:
                        img_e2e_res['result'][i][entity[0]] = []
                        img_e2e_res['result'][i][entity[0]].append(text[entity[1]:entity[2]])
    ...

最後に4.を実装します。条件分岐はreadtext()に記載されており、これを変更します。

ocr.py

def readtext(...):
    ...
    pp_result = None

    # Send args and models to the MMOCR model inference API
    # and call post-processing functions for the output
    if self.detect_model and self.recog_model and self.ner_model: # 追記
        det_recog_result = self.det_recog_ner_inference( # 追記
            self.detect_model, self.recog_model, ner_model=self.ner_model) # 追記
    elif self.detect_model and self.recog_model: # 変更
        det_recog_result = self.det_recog_kie_inference(
            self.detect_model, self.recog_model, kie_model=self.kie_model)
        pp_result = self.det_recog_pp(det_recog_result)
    else:
        for model in list(
                filter(None, [self.recog_model, self.detect_model])):
            result = self.single_inference(model, args.arrays,
                                           args.batch_mode,
                                           args.single_batch_size)
            pp_result = self.single_pp(result, model)

    return pp_result

実行

これらの変更を加えて、以下のコマンドによりOCR+NERを実行してみます。
また、本記事で学習させたwikiモデルを使用します。

python mmocr/utils/ocr.py demo/任意の画像 \
    --det PS_CTW \
    --recog SAR \
    --ner 適当な文字列 \
    --ner-config configs/ner/bert_softmax/bert_softmax_wiki_18e.py \
    --ner-ckpt work_dirs/wiki/epoch_18.pth

実行結果として入力画像に対するBBOX内のテキストとそれに対するNER結果が表示されます。

本記事について

本記事はTRUST SMITH株式会社でのインターン中に書きました。
TRUST SMITH株式会社は東大発のAIベンチャーです。
インターンやポスト、提供サービスに興味がおありでしたら以下から会社HPにアクセスしていただけます。
TRUST SMITH株式会社
 SMITH&VISION