# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np import os import sys __dir__ = os.path.dirname(os.path.abspath(__file__)) sys.path.append(__dir__) sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '..'))) os.environ["FLAGS_allocator_strategy"] = 'auto_growth' import cv2 import json import paddle import paddle.distributed as dist from ppocr.data import create_operators, transform from ppocr.modeling.architectures import build_model from ppocr.postprocess import build_post_process from ppocr.utils.save_load import load_model from ppocr.utils.visual import draw_re_results from ppocr.utils.logging import get_logger from ppocr.utils.utility import get_image_file_list, load_vqa_bio_label_maps, print_dict from tools.program import ArgsParser, load_config, merge_config from tools.infer_kie_token_ser import SerPredictor class ReArgsParser(ArgsParser): def __init__(self): super(ReArgsParser, self).__init__() self.add_argument( "-c_ser", "--config_ser", help="ser configuration file to use") self.add_argument( "-o_ser", "--opt_ser", nargs='+', help="set ser configuration options ") def parse_args(self, argv=None): args = super(ReArgsParser, self).parse_args(argv) assert args.config_ser is not None, \ "Please specify --config_ser=ser_configure_file_path." args.opt_ser = self._parse_opt(args.opt_ser) return args def make_input(ser_inputs, ser_results): entities_labels = {'HEADER': 0, 'QUESTION': 1, 'ANSWER': 2} entities = ser_inputs[8][0] ser_results = ser_results[0] assert len(entities) == len(ser_results) # entities start = [] end = [] label = [] entity_idx_dict = {} for i, (res, entity) in enumerate(zip(ser_results, entities)): if res['pred'] == 'O': continue entity_idx_dict[len(start)] = i start.append(entity['start']) end.append(entity['end']) label.append(entities_labels[res['pred']]) entities = dict(start=start, end=end, label=label) # relations head = [] tail = [] for i in range(len(entities["label"])): for j in range(len(entities["label"])): if entities["label"][i] == 1 and entities["label"][j] == 2: head.append(i) tail.append(j) relations = dict(head=head, tail=tail) batch_size = ser_inputs[0].shape[0] entities_batch = [] relations_batch = [] entity_idx_dict_batch = [] for b in range(batch_size): entities_batch.append(entities) relations_batch.append(relations) entity_idx_dict_batch.append(entity_idx_dict) ser_inputs[8] = entities_batch ser_inputs.append(relations_batch) # remove ocr_info segment_offset_id and label in ser input ser_inputs.pop(7) ser_inputs.pop(6) ser_inputs.pop(5) return ser_inputs, entity_idx_dict_batch class SerRePredictor(object): def __init__(self, config, ser_config): global_config = config['Global'] if "infer_mode" in global_config: ser_config["Global"]["infer_mode"] = global_config["infer_mode"] self.ser_engine = SerPredictor(ser_config) # init re model # build post process self.post_process_class = build_post_process(config['PostProcess'], global_config) # build model self.model = build_model(config['Architecture']) load_model( config, self.model, model_type=config['Architecture']["model_type"]) self.model.eval() def __call__(self, data): ser_results, ser_inputs = self.ser_engine(data) re_input, entity_idx_dict_batch = make_input(ser_inputs, ser_results) preds = self.model(re_input) post_result = self.post_process_class( preds, ser_results=ser_results, entity_idx_dict_batch=entity_idx_dict_batch) return post_result def preprocess(): FLAGS = ReArgsParser().parse_args() config = load_config(FLAGS.config) config = merge_config(config, FLAGS.opt) ser_config = load_config(FLAGS.config_ser) ser_config = merge_config(ser_config, FLAGS.opt_ser) logger = get_logger() # check if set use_gpu=True in paddlepaddle cpu version use_gpu = config['Global']['use_gpu'] device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu' device = paddle.set_device(device) logger.info('{} re config {}'.format('*' * 10, '*' * 10)) print_dict(config, logger) logger.info('\n') logger.info('{} ser config {}'.format('*' * 10, '*' * 10)) print_dict(ser_config, logger) logger.info('train with paddle {} and device {}'.format(paddle.__version__, device)) return config, ser_config, device, logger if __name__ == '__main__': config, ser_config, device, logger = preprocess() os.makedirs(config['Global']['save_res_path'], exist_ok=True) ser_re_engine = SerRePredictor(config, ser_config) if config["Global"].get("infer_mode", None) is False: data_dir = config['Eval']['dataset']['data_dir'] with open(config['Global']['infer_img'], "rb") as f: infer_imgs = f.readlines() else: infer_imgs = get_image_file_list(config['Global']['infer_img']) with open( os.path.join(config['Global']['save_res_path'], "infer_results.txt"), "w", encoding='utf-8') as fout: for idx, info in enumerate(infer_imgs): if config["Global"].get("infer_mode", None) is False: data_line = info.decode('utf-8') substr = data_line.strip("\n").split("\t") img_path = os.path.join(data_dir, substr[0]) data = {'img_path': img_path, 'label': substr[1]} else: img_path = info data = {'img_path': img_path} save_img_path = os.path.join( config['Global']['save_res_path'], os.path.splitext(os.path.basename(img_path))[0] + "_ser_re.jpg") result = ser_re_engine(data) result = result[0] fout.write(img_path + "\t" + json.dumps( result, ensure_ascii=False) + "\n") img_res = draw_re_results(img_path, result) cv2.imwrite(save_img_path, img_res) logger.info("process: [{}/{}], save result to {}".format( idx, len(infer_imgs), save_img_path))