# 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 glob
import os
import pathlib
import sys
import time
import xml
import xml.etree.ElementTree as ET
from . import utils
from .. import procs
# Add import path for src modules
currentdir = pathlib.Path(__file__).resolve().parent
sys.path.append(str(currentdir) + "/../../src/separate_pages_ssd")
sys.path.append(str(currentdir) + "/../../src/ndl_layout")
sys.path.append(str(currentdir) + "/../../src/deskew_HT")
sys.path.append(str(currentdir) + "/../../src/text_recognition")
# supported image type list
supported_img_ext = ['.jpg', '.jpeg', '.jp2']
class OcrInferencer:
"""
推論実行時の関数や推論の設定値を保持します。
Attributes
----------
full_proc_list : list
全推論処理のリストです。
proc_list : list
本実行処理における推論処理のリストです。
cfg : dict
本実行処理における設定情報です。
"""
def __init__(self, cfg):
"""
Parameters
----------
cfg : dict
本実行処理における設定情報です。
"""
# inference process class list in order
self.full_proc_list = [
procs.PageSeparation, # 0: ノド元分割 出力：（画像：あり、XML：なし、TXT：なし）
procs.PageDeskewProcess, # 1: 傾き補正 出力：（画像：あり、XML：なし、TXT：なし）
procs.LayoutExtractionProcess, # 2: レイアウト抽出 出力：（画像：あり、XML：あり、TXT：なし）
procs.LineOcrProcess, # 3: 文字認識(OCR) 出力：（画像：あり、XML：あり、TXT：あり）
]
self.proc_list = self._create_proc_list(cfg)
self.cfg = cfg
self.time_statistics = []
self.xml_template = '<?xml version="1.0" encoding="utf-8" standalone="yes"?>\n<OCRDATASET></OCRDATASET>'
def run(self):
"""
self.cfgに保存された設定に基づいた推論処理を実行します。
"""
if len(self.cfg['input_dirs']) == 0:
print('[ERROR] Input directory list is empty', file=sys.stderr)
return
# input dir loop
for input_dir in self.cfg['input_dirs']:
if self.cfg['input_structure'] in ['t']:
single_outputdir_data_list = self._get_single_dir_data_from_tosho_data(input_dir)
else:
single_outputdir_data_list = self._get_single_dir_data(input_dir)
if single_outputdir_data_list is None:
print('[ERROR] Input data list is empty', file=sys.stderr)
continue
print(single_outputdir_data_list)
# do infer with input data for single output data dir
for single_outputdir_data in single_outputdir_data_list:
print(single_outputdir_data)
if single_outputdir_data is None:
continue
pred_list = self._infer(single_outputdir_data)
# save inferenced xml in xml directory
if (self.cfg['save_xml'] or self.cfg['partial_infer']) and (self.cfg['proc_range']['end'] > 1):
self._save_pred_xml(single_outputdir_data['output_dir'], [single_data['xml'] for single_data in pred_list])
if len(self.time_statistics) == 0:
print('================== NO VALID INFERENCE ==================')
else:
average = sum(self.time_statistics) / len(self.time_statistics)
print('================== PROCESSING TIME ==================')
print('Average processing time : {0} sec / image file '.format(average))
return
def _infer(self, single_outputdir_data):
"""
self.cfgに保存された設定に基づき、XML一つ分のデータに対する推論処理を実行します。
Parameters
----------
single_outputdir_data : dict
XML一つ分のデータ（基本的に1書籍分を想定）の入力データ情報。
画像ファイルパスのリスト、それらに対応するXMLデータを含みます。
Returns
-------
pred_list : list
1ページ分の推論結果を要素に持つ推論結果のリスト。
各結果は辞書型で保持されています。
"""
# single_outputdir_data dictionary include [key, value] pairs as below
# (xml is not always included)
# [key, value]: ['img', numpy.ndarray], ['xml', xml_tree]
pred_list = []
pred_xml_dict_for_dump = {}
if self.cfg['dump']:
dump_dir = os.path.join(single_outputdir_data['output_dir'], 'dump')
os.makedirs(dump_dir, exist_ok=True)
for proc in self.proc_list:
pred_xml_dict_for_dump[proc.proc_name] = []
proc_dump_dir = os.path.join(dump_dir, proc.proc_name)
os.makedirs(proc_dump_dir, exist_ok=True)
for img_path in single_outputdir_data['img_list']:
single_image_file_data = self._get_single_image_file_data(img_path, single_outputdir_data)
output_dir = single_outputdir_data['output_dir']
if single_image_file_data is None:
print('[ERROR] Failed to get single page input data for image:{0}'.format(img_path), file=sys.stderr)
continue
print('######## START PAGE INFERENCE PROCESS ########')
start_page = time.time()
for proc in self.proc_list:
single_page_output = []
for idx, single_data_input in enumerate(single_image_file_data):
single_data_output = proc.do(idx, single_data_input)
single_page_output.extend(single_data_output)
# save inference result data to dump
if self.cfg['dump'] and 'xml' in single_image_file_data[0].keys():
pred_xml_dict_for_dump[proc.proc_name].append(single_image_file_data[0]['xml'])
single_image_file_data = single_page_output
single_image_file_output = single_image_file_data
self.time_statistics.append(time.time() - start_page)
if self.cfg['save_image'] or self.cfg['partial_infer']:
# save inferenced result drawn image in pred_img directory
for single_data_output in single_image_file_output:
# save input image while partial inference
if self.cfg['partial_infer']:
img_output_dir = os.path.join(output_dir, 'img')
self._save_image(single_data_output['img'], single_data_output['img_file_name'], img_output_dir)
pred_img = self._create_result_image(single_data_output, self.proc_list[-1].proc_name)
img_output_dir = os.path.join(output_dir, 'pred_img')
self._save_image(pred_img, single_data_output['img_file_name'], img_output_dir)
# save inferenced result text for this page
if self.cfg['proc_range']['end'] > 2:
sum_main_txt = ''
sum_cap_txt = ''
for single_data_output in single_image_file_output:
main_txt, cap_txt = self._create_result_txt(single_data_output['xml'])
sum_main_txt += main_txt + '\n'
sum_cap_txt += sum_cap_txt + '\n'
self._save_pred_txt(sum_main_txt, sum_cap_txt, os.path.basename(img_path), single_outputdir_data['output_dir'])
# add inference result for single image file data to pred_list, including XML data
pred_list.extend(single_image_file_output)
print('######## END PAGE INFERENCE PROCESS ########')
return pred_list
def _get_single_dir_data(self, input_dir):
"""
XML一つ分の入力データに関する情報を整理して取得します。
Parameters
----------
input_dir : str
XML一つ分の入力データが保存されているディレクトリパスです。
Returns
-------
# Fixme
single_dir_data : dict
XML一つ分のデータ（基本的に1PID分を想定）の入力データ情報です。
画像ファイルパスのリスト、それらに対応するXMLデータを含みます。
"""
single_dir_data = {'input_dir': os.path.abspath(input_dir)}
single_dir_data['img_list'] = []
# get img list of input directory
if self.cfg['input_structure'] in ['w']:
for ext in supported_img_ext:
single_dir_data['img_list'].extend(sorted(glob.glob(os.path.join(input_dir, '*{0}'.format(ext)))))
elif self.cfg['input_structure'] in ['f']:
stem, ext = os.path.splitext(os.path.basename(input_dir))
if ext in supported_img_ext:
single_dir_data['img_list'] = [input_dir]
else:
print('[ERROR] This file is not supported type : {0}'.format(input_dir), file=sys.stderr)
elif not os.path.isdir(os.path.join(input_dir, 'img')):
print('[ERROR] Input img diretctory not found in {}'.format(input_dir), file=sys.stderr)
return None
else:
for ext in supported_img_ext:
single_dir_data['img_list'].extend(sorted(glob.glob(os.path.join(input_dir, 'img/*{0}'.format(ext)))))
# check xml file number and load xml data if needed
if self.cfg['proc_range']['start'] > 2:
if self.cfg['input_structure'] in ['f']:
print('[ERROR] Single image file input mode does not support partial inference wich need xml file input.', file=sys.stderr)
return None
input_xml = None
xml_file_list = glob.glob(os.path.join(input_dir, 'xml/*.xml'))
if len(xml_file_list) > 1:
print('[ERROR] Input xml file must be only one, but there is {0} xml files in {1}.'.format(
len(xml_file_list), os.path.join(self.cfg['input_root'], 'xml')), file=sys.stderr)
return None
elif len(xml_file_list) == 0:
print('[ERROR] There is no input xml files in {0}.'.format(os.path.join(input_dir, 'xml')), file=sys.stderr)
return None
else:
input_xml = xml_file_list[0]
try:
single_dir_data['xml'] = ET.parse(input_xml)
except xml.etree.ElementTree.ParseError as err:
print("[ERROR] XML parse error : {0}".format(input_xml), file=sys.stderr)
return None
# prepare output dir for inferensce result with this input dir
if self.cfg['input_structure'] in ['f']:
stem, ext = os.path.splitext(os.path.basename(input_dir))
output_dir = os.path.join(self.cfg['output_root'], stem)
elif self.cfg['input_structure'] in ['i', 's']:
dir_name = os.path.basename(input_dir)
output_dir = os.path.join(self.cfg['output_root'], dir_name)
elif self.cfg['input_structure'] in ['w']:
input_dir_names = input_dir.split('/')
dir_name = input_dir_names[-3][0] + input_dir_names[-2] + input_dir_names[-1]
output_dir = os.path.join(self.cfg['output_root'], dir_name)
else:
print('[ERROR] Unexpected input directory structure type: {}.'.format(self.cfg['input_structure']), file=sys.stderr)
return None
# output directory existance check
output_dir = utils.mkdir_with_duplication_check(output_dir)
single_dir_data['output_dir'] = output_dir
return [single_dir_data]
def _get_single_dir_data_from_tosho_data(self, input_dir):
"""
XML一つ分の入力データに関する情報を整理して取得します。
Parameters
----------
input_dir : str
tosho data形式のセクションごとのディレクトリパスです。
Returns
-------
single_dir_data_list : list
XML一つ分のデータ（基本的に1PID分を想定）の入力データ情報のリストです。
1つの要素に画像ファイルパスのリスト、それらに対応するXMLデータを含みます。
"""
single_dir_data_list = []
# get img list of input directory
tmp_img_list = sorted(glob.glob(os.path.join(input_dir, '*.jp2')))
tmp_img_list.extend(sorted(glob.glob(os.path.join(input_dir, '*.jpg'))))
pid_list = []
for img in tmp_img_list:
pid = os.path.basename(img).split('_')[0]
if pid not in pid_list:
pid_list.append(pid)
for pid in pid_list:
single_dir_data = {'input_dir': os.path.abspath(input_dir),
'img_list': [img for img in tmp_img_list if os.path.basename(img).startswith(pid)]}
# prepare output dir for inferensce result with this input dir
output_dir = os.path.join(self.cfg['output_root'], pid)
# output directory existance check
os.makedirs(output_dir, exist_ok=True)
single_dir_data['output_dir'] = output_dir
single_dir_data_list.append(single_dir_data)
return single_dir_data_list
def _get_single_image_file_data(self, img_path, single_dir_data):
"""
1ページ分の入力データに関する情報を整理して取得します。
Parameters
----------
img_path : str
入力画像データのパスです。
single_dir_data : dict
1書籍分の入力データに関する情報を保持する辞書型データです。
xmlファイルへのパス、結果を出力するディレクトリのパスなどを含みます。
Returns
-------
single_image_file_data : dict
1ページ分のデータの入力データ情報です。
画像ファイルのパスとnumpy.ndarray形式の画像データ、その画像に対応するXMLデータを含みます。
"""
single_image_file_data = [{
'img_path': img_path,
'img_file_name': os.path.basename(img_path),
'output_dir': single_dir_data['output_dir']
}]
full_xml = None
if 'xml' in single_dir_data.keys():
full_xml = single_dir_data['xml']
# get img data for single page
orig_img = cv2.imread(img_path)
if orig_img is None:
print('[ERROR] Image read error : {0}'.format(img_path), file=sys.stderr)
return None
single_image_file_data[0]['img'] = orig_img
# return if this proc needs only img data for input
if full_xml is None:
return single_image_file_data
# get xml data for single page
image_name = os.path.basename(img_path)
for page in full_xml.getroot().iter('PAGE'):
if page.attrib['IMAGENAME'] == image_name:
node = ET.fromstring(self.xml_template)
node.append(page)
tree = ET.ElementTree(node)
single_image_file_data[0]['xml'] = tree
break
# [TODO] 画像データに対応するXMLデータが見つからなかった場合の対応
if 'xml' not in single_image_file_data[0].keys():
print('[ERROR] Input XML data for page {} not found.'.format(img_path), file=sys.stderr)
return single_image_file_data
def _create_proc_list(self, cfg):
"""
推論の設定情報に基づき、実行する推論処理のリストを作成します。
Parameters
----------
cfg : dict
推論実行時の設定情報を保存した辞書型データ。
"""
proc_list = []
for i in range(cfg['proc_range']['start'], cfg['proc_range']['end'] + 1):
proc_list.append(self.full_proc_list[i](cfg, i))
return proc_list
def _save_pred_xml(self, output_dir, pred_list):
"""
推論結果のXMLデータをまとめたXMLファイルを生成して保存します。
Parameters
----------
output_dir : str
推論結果を保存するディレクトリのパスです。
pred_list : list
1ページ分の推論結果を要素に持つ推論結果のリスト。
各結果は辞書型で保持されています。
"""
xml_dir = os.path.join(output_dir, 'xml')
os.makedirs(xml_dir, exist_ok=True)
# basically, output_dir is supposed to be PID, so it used as xml filename
xml_path = os.path.join(xml_dir, '{}.xml'.format(os.path.basename(output_dir)))
pred_xml = self._parse_pred_list_to_save(pred_list)
utils.save_xml(pred_xml, xml_path)
return
def _save_image(self, pred_img, orig_img_name, img_output_dir, id=''):
"""
指定されたディレクトリに画像データを保存します。
画像データは入力に使用したものと推論結果を重畳したものの２種類が想定されています。
Parameters
----------
pred_img : numpy.ndarray
保存する画像データ。
orig_img_name : str
もともとの入力画像のファイル名。
基本的にはこのファイル名と同名で保存します。
img_output_dir : str
画像ファイルの保存先のディレクトリパス。
id : str
もともとの入力画像のファイル名に追加する処理結果ごとのidです。
一つの入力画像から複数の画像データが出力される処理がある場合に必要になります。
"""
os.makedirs(img_output_dir, exist_ok=True)
stem, ext = os.path.splitext(orig_img_name)
orig_img_name = stem + '.jpg'
if id != '':
stem, ext = os.path.splitext(orig_img_name)
orig_img_name = stem + '_' + id + ext
img_path = os.path.join(img_output_dir, orig_img_name)
try:
cv2.imwrite(img_path, pred_img)
except OSError as err:
print("[ERROR] Image save error: {0}".format(err), file=sys.stderr)
raise OSError
return
def _save_pred_txt(self, main_txt, cap_txt, orig_img_name, output_dir):
"""
指定されたディレクトリに推論結果のテキストデータを保存します。
Parameters
----------
main_txt : str
本文＋キャプションの推論結果のテキストデータです
cap_txt : str
キャプションのみの推論結果のテキストデータです
orig_img_name : str
もともとの入力画像ファイル名。
基本的にはこのファイル名と同名で保存します。
img_output_dir : str
画像ファイルの保存先のディレクトリパス。
"""
txt_dir = os.path.join(output_dir, 'txt')
os.makedirs(txt_dir, exist_ok=True)
stem, _ = os.path.splitext(orig_img_name)
txt_path = os.path.join(txt_dir, stem + '_cap.txt')
try:
with open(txt_path, 'w') as f:
f.write(cap_txt)
except OSError as err:
print("[ERROR] Caption text save error: {0}".format(err), file=sys.stderr)
raise OSError
stem, _ = os.path.splitext(orig_img_name)
txt_path = os.path.join(txt_dir, stem + '_main.txt')
try:
with open(txt_path, 'w') as f:
f.write(main_txt)
except OSError as err:
print("[ERROR] Main text save error: {0}".format(err), file=sys.stderr)
raise OSError
return
def _parse_pred_list_to_save(self, pred_list):
"""
推論結果のXMLを要素に持つリストから、ファイルに保存するための一つのXMLデータを生成します。
Parameters
----------
pred_list : list
推論結果のXMLを要素に持つリスト。
"""
ET.register_namespace('', 'NDLOCRDATASET')
node = ET.fromstring(self.xml_template)
for single_xml_tree in pred_list:
root = single_xml_tree.getroot()
for element in root:
node.append(element)
tree = ET.ElementTree(node)
return tree
def _create_result_image(self, result, proc_name):
"""
推論結果を入力画像に重畳した画像データを生成します。
Parameters
----------
result : dict
1ページ分の推論結果を持つ辞書型データ。
proc_name : str
重畳を行う結果を出力した推論処理の名前。
"""
if 'dump_img' in result.keys():
dump_img = copy.deepcopy(result['dump_img'])
else:
dump_img = copy.deepcopy(result['img'])
if 'xml' in result.keys() and result['xml'] is not None:
# draw inference result on input image
cv2.putText(dump_img, proc_name, (0, 50),
cv2.FONT_HERSHEY_PLAIN, 4, (0, 0, 0), 5, cv2.LINE_AA)
pass
else:
cv2.putText(dump_img, proc_name, (0, 50),
cv2.FONT_HERSHEY_PLAIN, 4, (0, 0, 0), 5, cv2.LINE_AA)
return dump_img
def _create_result_txt(self, xml_data):
"""
推論結果のxmlデータからテキストデータを生成します。
Parameters
----------
xml_data :
1ページ分の推論結果を持つxmlデータ。
"""
main_txt = ''
cap_txt = ''
for page_xml in xml_data.iter('PAGE'):
for line_xml in page_xml.iter('LINE'):
main_txt += line_xml.attrib['STRING']
main_txt += '\n'
if line_xml.attrib['TYPE'] == 'キャプション':
cap_txt += line_xml.attrib['STRING']
cap_txt += '\n'
return main_txt, cap_txt