cli/procs/base_proc.py

# 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