# Copyright (c) 2022, National Diet Library, Japan
#
# This software is released under the CC BY 4.0.
# https://creativecommons.org/licenses/by/4.0/
import copy
import cv2
import os
class BaseInferenceProcess:
"""
各推論処理を実行するプロセスクラスを作るためのメタクラス。
Attributes
----------
proc_name : str
推論処理を実行するインスタンスが持つプロセス名。
[実行される順序を表す数字＋クラスごとの処理名]で構成されます。
cfg : dict
本推論実行における設定情報です。
"""
def __init__(self, cfg, pid, proc_type='_base_prep'):
"""
Parameters
----------
cfg : dict
本実行処理における設定情報です。
pid : int
実行される順序を表す数値。
proc_type : str
クラスごとに定義されている処理名。
"""
self.proc_name = str(pid) + proc_type
if not self._is_valid_cfg(cfg):
raise ValueError('Configuration validation error.')
else:
self.cfg = cfg
self.process_dump_dir = None
return True
def do(self, data_idx, input_data):
"""
推論処理を実行する際にOcrInferencerクラスから呼び出される推論実行関数。
入力データのバリデーションや推論処理、推論結果の保存などが含まれます。
本処理は基本的に継承先では変更されないことを想定しています。
Parameters
----------
data_idx : int
入力データのインデックス。
画像ファイル１つごとに入力データのリストが構成されます。
input_data : dict
推論処理を実行すつ対象の入力データ。
Returns
-------
result : dict
推論処理の結果を保持する辞書型データ。
基本的にinput_dataと同じ構造です。
"""
# input data valudation check
if not self._is_valid_input(input_data):
raise ValueError('Input data validation error.')
# run main inference process
result = self._run_process(input_data)
if result is None:
raise ValueError('Inference output error in {0}.'.format(self.proc_name))
# dump inference result
if self.cfg['dump']:
self._dump_result(input_data, result, data_idx)
return result
def _run_process(self, input_data):
"""
推論処理の本体部分。
処理内容は継承先のクラスで実装されることを想定しています。
Parameters
----------
input_data : dict
推論処理を実行する対象の入力データ。
Returns
-------
result : dict
推論処理の結果を保持する辞書型データ。
基本的にinput_dataと同じ構造です。
"""
print('### Base Inference Process ###')
result = copy.deepcopy(input_data)
return result
def _is_valid_cfg(self, cfg):
"""
推論処理全体の設定情報ではなく、クラス単位の設定情報に対するバリデーション。
バリデーションの内容は継承先のクラスで実装されることを想定しています。
Parameters
----------
cfg : dict
本推論実行における設定情報です。
Returns
-------
[変数なし] : bool
設定情報が正しければTrue, そうでなければFalseを返します。
"""
if cfg is None:
print('Given configuration data is None.')
return False
return True
def _is_valid_input(self, input_data):
"""
本クラスの推論処理における入力データのバリデーション。
バリデーションの内容は継承先のクラスで実装されることを想定しています。
Parameters
----------
input_data : dict
推論処理を実行する対象の入力データ。
Returns
-------
[変数なし] : bool
入力データが正しければTrue, そうでなければFalseを返します。
"""
return True
def _dump_result(self, input_data, result, data_idx):
"""
本クラスの推論処理結果をファイルに保存します。
dumpフラグが有効の場合にのみ実行されます。
Parameters
----------
input_data : dict
推論処理に利用した入力データ。
result : list
推論処理の結果を保持するリスト型データ。
各要素は基本的にinput_dataと同じ構造の辞書型データです。
data_idx : int
入力データのインデックス。
画像ファイル１つごとに入力データのリストが構成されます。
"""
self.process_dump_dir = os.path.join(os.path.join(input_data['output_dir'], 'dump'), self.proc_name)
for i, single_result in enumerate(result):
if 'img' in single_result.keys() and single_result['img'] is not None:
dump_img_name = os.path.basename(input_data['img_path']).split('.')[0] + '_' + str(data_idx) + '_' + str(i) + '.jpg'
self._dump_img_result(single_result, input_data['output_dir'], dump_img_name)
if 'xml' in single_result.keys() and single_result['xml'] is not None:
dump_xml_name = os.path.basename(input_data['img_path']).split('.')[0] + '_' + str(data_idx) + '_' + str(i) + '.xml'
self._dump_xml_result(single_result, input_data['output_dir'], dump_xml_name)
if 'txt' in single_result.keys() and single_result['txt'] is not None:
dump_txt_name = os.path.basename(input_data['img_path']).split('.')[0] + '_' + str(data_idx) + '_' + str(i) + '.txt'
self._dump_txt_result(single_result, input_data['output_dir'], dump_txt_name)
return
def _dump_img_result(self, single_result, output_dir, img_name):
"""
本クラスの推論処理結果(画像)をファイルに保存します。
dumpフラグが有効の場合にのみ実行されます。
Parameters
----------
single_result : dict
推論処理の結果を保持する辞書型データ。
output_dir : str
推論結果が保存されるディレクトリのパス。
img_name : str
入力データの画像ファイル名。
dumpされる画像ファイルのファイル名は入力のファイル名と同名(複数ある場合は連番を付与)となります。
"""
pred_img_dir = os.path.join(self.process_dump_dir, 'pred_img')
os.makedirs(pred_img_dir, exist_ok=True)
image_file_path = os.path.join(pred_img_dir, img_name)
dump_image = self._create_result_image(single_result)
try:
cv2.imwrite(image_file_path, dump_image)
except OSError as err:
print("Dump image save error: {0}".format(err))
raise OSError
return
def _dump_xml_result(self, single_result, output_dir, img_name):
"""
本クラスの推論処理結果(XML)をファイルに保存します。
dumpフラグが有効の場合にのみ実行されます。
Parameters
----------
single_result : dict
推論処理の結果を保持する辞書型データ。
output_dir : str
推論結果が保存されるディレクトリのパス。
img_name : str
入力データの画像ファイル名。
dumpされるXMLファイルのファイル名は入力のファイル名とほぼ同名（拡張子の変更、サフィックスや連番の追加のみ）となります。
"""
xml_dir = os.path.join(self.process_dump_dir, 'xml')
os.makedirs(xml_dir, exist_ok=True)
trum, _ = os.path.splitext(img_name)
xml_path = os.path.join(xml_dir, trum + '.xml')
try:
single_result['xml'].write(xml_path, encoding='utf-8', xml_declaration=True)
except OSError as err:
print("Dump xml save error: {0}".format(err))
raise OSError
return
def _dump_txt_result(self, single_result, output_dir, img_name):
"""
本クラスの推論処理結果(テキスト)をファイルに保存します。
dumpフラグが有効の場合にのみ実行されます。
Parameters
----------
single_result : dict
推論処理の結果を保持する辞書型データ。
output_dir : str
推論結果が保存されるディレクトリのパス。
img_name : str
入力データの画像ファイル名。
dumpされるテキストファイルのファイル名は入力のファイル名とほぼ同名（拡張子の変更、サフィックスや連番の追加のみ）となります。
"""
txt_dir = os.path.join(self.process_dump_dir, 'txt')
os.makedirs(txt_dir, exist_ok=True)
trum, _ = os.path.splitext(img_name)
txt_path = os.path.join(txt_dir, trum + '_main.txt')
try:
with open(txt_path, 'w') as f:
f.write(single_result['txt'])
except OSError as err:
print("Dump text save error: {0}".format(err))
raise OSError
return
def _create_result_image(self, single_result):
"""
推論結果を入力の画像に重畳した画像データを生成します。
Parameters
----------
single_result : dict
推論処理の結果を保持する辞書型データ。
"""
dump_img = None
if 'dump_img' in single_result.keys():
dump_img = copy.deepcopy(single_result['dump_img'])
else:
dump_img = copy.deepcopy(single_result['img'])
if 'xml' in single_result.keys() and single_result['xml'] is not None:
# draw single inferenceresult on input image
# this should be implemeted in each child class
cv2.putText(dump_img, 'dump' + self.proc_name, (0, 50),
cv2.FONT_HERSHEY_PLAIN, 4, (255, 0, 0), 5, cv2.LINE_AA)
pass
else:
cv2.putText(dump_img, 'dump' + self.proc_name, (0, 50),
cv2.FONT_HERSHEY_PLAIN, 4, (255, 255, 0), 5, cv2.LINE_AA)
return dump_img