OpenMMLab の MMOCR に日本語の学習をさせたい【dev-1.x】

公開日：2022年10月27日
最終更新日：2022年11月24日

はじめに

前回「MMOCR==0.6.2」を使って日本語の学習を行いました。
touch-sp.hatenablog.com
今回は「MMOCR==1.0.0rc2」を使ってみたいと思います。

PC環境

前回と一緒です。

Windows 11
CUDA 11.6.2
Python 3.10.8

Python環境構築

pip install torch==1.12.1 torchvision==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install openmim==0.3.2
pip install mmengine==0.2.0
pip install mmcv==2.0.0rc1 -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.12.0/index.html
pip install mmdet==3.0.0rc2
pip install mmocr==1.0.0rc2
pip install openpyxl==3.0.10
pip install trdg==1.8.0

モデルによっては以下が必要になります。

pip install albumentations==1.3.0

「mmocr==1.0.0rc3」も試しましたが問題なく動きました。

「mmcv==2.0.0rc2」にすると動きませんでした。

学習データの作成

前回同様「TextRecognitionDataGenerator」を使わせて頂きました。

前回と異なる点は「TextRecognitionDataGenerator」が出力するtextファイルをjsonファイルに変更する必要があることです。

以下のスクリプトで変換が可能です。

import argparse
import json

parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, help='text file')
parser.add_argument('--output', type = str, help='json file')
args = parser.parse_args()

input_fname = args.input
output_fname = args.output

with open(input_fname, 'r', encoding='utf-8') as f:
    lines = f.readlines()

lines = [x.strip() for x in lines]

data_list = []
for line in lines:
    img_path, text = line.split(' ')
    data = {
        'img_path': img_path,
        'instances':[{'text':text}]
        }
    data_list.append(data)

result = {
    'metainfo':{
        'dataset_type':'TextRecogDataset',
        'task_name':'textrecog'
    },
    'data_list':data_list
}

with open(output_fname, 'w', encoding='utf-8') as f:
    json.dump(result, f, indent=2, ensure_ascii=False)

Configファイルを作成する

train = dict(
    type='OCRDataset',
    data_prefix=dict(img_path='mnt/ramdisk/max/90kDICT32px'),
    ann_file='train_labels.json',
    test_mode=False,
    pipeline=None)

test = dict(
    type='OCRDataset',
    data_prefix=dict(img_path='mnt/ramdisk/max/90kDICT32px'),
    ann_file='test_labels.json',
    test_mode=True,
    pipeline=None)

default_scope = 'mmocr'

env_cfg = dict(
    cudnn_benchmark=True,
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
    dist_cfg=dict(backend='nccl'))

randomness = dict(seed=None)

default_hooks = dict(
    timer=dict(type='IterTimerHook'),
    logger=dict(type='LoggerHook', interval=100),
    param_scheduler=dict(type='ParamSchedulerHook'),
    checkpoint=dict(type='CheckpointHook', interval=1),
    sampler_seed=dict(type='DistSamplerSeedHook'),
    sync_buffer=dict(type='SyncBuffersHook'),
    visualization=dict(
        type='VisualizationHook',
        interval=1,
        enable=False,
        show=False,
        draw_gt=False,
        draw_pred=False))

log_level = 'INFO'
log_processor = dict(type='LogProcessor', window_size=10, by_epoch=True)
load_from = None
resume = False

val_evaluator = dict(
    type='Evaluator',
    metrics=[
        dict(
            type='WordMetric',
            mode=['exact', 'ignore_case', 'ignore_case_symbol']),
        dict(type='CharMetric')
    ])

test_evaluator = dict(
    type='Evaluator',
    metrics=[
        dict(
            type='WordMetric',
            mode=['exact', 'ignore_case', 'ignore_case_symbol']),
        dict(type='CharMetric')
    ])

vis_backends = [dict(type='LocalVisBackend')]

visualizer = dict(
    type='TextRecogLocalVisualizer',
    name='visualizer',
    vis_backends=[dict(type='LocalVisBackend')])

optim_wrapper = dict(
    type='OptimWrapper', optimizer=dict(type='Adam', lr=0.0003))
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=5, val_interval=1)

val_cfg = dict(type='ValLoop')
test_cfg = dict(type='TestLoop')

param_scheduler = [dict(type='MultiStepLR', milestones=[3, 4], end=5)]

file_client_args = dict(backend='disk')

dictionary = dict(
    type='Dictionary',
    dict_file=
    'd:/python/preocrenv/lib/site-packages/mmocr/.mim/configs/textrecog/satrn/../../../dicts/english_digits_symbols.txt',
    with_padding=True,
    with_unknown=True,
    same_start_end=True,
    with_start=True,
    with_end=True)

model = dict(
    type='SATRN',
    backbone=dict(type='ShallowCNN', input_channels=3, hidden_dim=512),
    encoder=dict(
        type='SATRNEncoder',
        n_layers=12,
        n_head=8,
        d_k=64,
        d_v=64,
        d_model=512,
        n_position=100,
        d_inner=2048,
        dropout=0.1),
    decoder=dict(
        type='NRTRDecoder',
        n_layers=6,
        d_embedding=512,
        n_head=8,
        d_model=512,
        d_inner=2048,
        d_k=64,
        d_v=64,
        module_loss=dict(
            type='CEModuleLoss', flatten=True, ignore_first_char=True),
        dictionary=dict(
            type='Dictionary',
            dict_file=
            'd:/python/preocrenv/lib/site-packages/mmocr/.mim/configs/textrecog/satrn/../../../dicts/english_digits_symbols.txt',
            with_padding=True,
            with_unknown=True,
            same_start_end=True,
            with_start=True,
            with_end=True),
        max_seq_len=25,
        postprocessor=dict(type='AttentionPostprocessor')),
    data_preprocessor=dict(
        type='TextRecogDataPreprocessor',
        mean=[123.675, 116.28, 103.53],
        std=[58.395, 57.12, 57.375]))

train_pipeline = [
    dict(
        type='LoadImageFromFile',
        file_client_args=dict(backend='disk'),
        ignore_empty=True,
        min_size=2),
    dict(type='LoadOCRAnnotations', with_text=True),
    dict(type='Resize', scale=(100, 32), keep_ratio=False),
    dict(
        type='PackTextRecogInputs',
        meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
]

test_pipeline = [
    dict(type='LoadImageFromFile', file_client_args=dict(backend='disk')),
    dict(type='Resize', scale=(100, 32), keep_ratio=False),
    dict(type='LoadOCRAnnotations', with_text=True),
    dict(
        type='PackTextRecogInputs',
        meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
]

train_dataset = dict(
    type='ConcatDataset',
    datasets=[train],
    pipeline=[
        dict(
            type='LoadImageFromFile',
            file_client_args=dict(backend='disk'),
            ignore_empty=True,
            min_size=2),
        dict(type='LoadOCRAnnotations', with_text=True),
        dict(type='Resize', scale=(100, 32), keep_ratio=False),
        dict(
            type='PackTextRecogInputs',
            meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
    ])

test_dataset = dict(
    type='ConcatDataset',
    datasets=[test],
    pipeline=[
        dict(type='LoadImageFromFile', file_client_args=dict(backend='disk')),
        dict(type='Resize', scale=(100, 32), keep_ratio=False),
        dict(type='LoadOCRAnnotations', with_text=True),
        dict(
            type='PackTextRecogInputs',
            meta_keys=('img_path', 'ori_shape', 'img_shape', 'valid_ratio'))
    ])

train_dataloader = dict(
    batch_size=128,
    num_workers=24,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    dataset=train_dataset)

test_dataloader = dict(
    batch_size=1,
    num_workers=4,
    persistent_workers=True,
    drop_last=False,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=test_dataset)

val_dataloader = dict(
    batch_size=1,
    num_workers=4,
    persistent_workers=True,
    drop_last=False,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=test_dataset)

auto_scale_lr = dict(base_batch_size=512)

学習用ファイルを実行する

import os
from mmengine.config import Config
from mmengine.registry import RUNNERS
from mmengine.runner import Runner

def main():
    cfg = Config.fromfile('satrn_japanese_cfg.py')

    os.makedirs('satrn_output', exist_ok=True)

    ####
    ## modify configuration file
    ####

    # Set output dir
    cfg.work_dir = 'satrn_output'

    # Path to dictionary file
    cfg.dictionary.dict_file = 'dicts.txt'
    cfg.model.decoder.dictionary = cfg.dictionary
    
    # Path to annotation file and image folder
    cfg.train.data_prefix.img_path = 'train'
    cfg.train.ann_file = 'train_labels.json'
    cfg.train_dataset.datasets = [cfg.train]
    cfg.train_dataloader.dataset = cfg.train_dataset

    cfg.test.data_prefix.img_path = 'test'
    cfg.test.ann_file = 'test_labels.json'
    cfg.test_dataset.datasets = [cfg.test]
    cfg.test_dataloader.dataset = cfg.test_dataset

    cfg.val_dataloader.dataset = cfg.test_dataset

    # Modify cuda setting
    cfg.gpu_ids = range(1)
    cfg.device = 'cuda'

    # Others
    cfg.train_dataloader.batch_size = 32
    cfg.train_dataloader.num_workers = 8
    cfg.model.decoder.max_seq_len = 35
    cfg.train_cfg.max_epochs = 1 # default 5 
    
    # Build the runner from config
    if 'runner_type' not in cfg:
        # build the default runner
        runner = Runner.from_cfg(cfg)
    else:
        # build customized runner from the registry
        # if 'runner_type' is set in the cfg
        runner = RUNNERS.build(cfg)

    # Start training
    runner.train()
    
if __name__ == '__main__':
    main()

補足①

上記のConfigファイルを見ると画像を(100, 32)にリサイズしています。

dict(type='Resize', scale=(100, 32), keep_ratio=False),

scale=(200, 32)に変更するとmodel.encoder.n_positionも200に変更する必要があります。

補足②

新しい記事も書いています。そちらも読んでみて下さい。
touch-sp.hatenablog.com

その他

スクリプトはGitHubで公開しています。
github.com