diff --git a/ppstructure/table/predict_structure.py b/ppstructure/table/predict_structure.py index 0bf100852b9e9d501dfc858d8ce0787da42a61ed..08e381a846f1e8b4d38918e1031f5b219fed54e2 100755 --- a/ppstructure/table/predict_structure.py +++ b/ppstructure/table/predict_structure.py @@ -68,6 +68,7 @@ def build_pre_process_list(args): class TableStructurer(object): def __init__(self, args): + self.args = args self.use_onnx = args.use_onnx pre_process_list = build_pre_process_list(args) if args.table_algorithm not in ['TableMaster']: @@ -89,8 +90,31 @@ class TableStructurer(object): self.predictor, self.input_tensor, self.output_tensors, self.config = \ utility.create_predictor(args, 'table', logger) + if args.benchmark: + import auto_log + pid = os.getpid() + gpu_id = utility.get_infer_gpuid() + self.autolog = auto_log.AutoLogger( + model_name="table", + model_precision=args.precision, + batch_size=1, + data_shape="dynamic", + save_path=None, #args.save_log_path, + inference_config=self.config, + pids=pid, + process_name=None, + gpu_ids=gpu_id if args.use_gpu else None, + time_keys=[ + 'preprocess_time', 'inference_time', 'postprocess_time' + ], + warmup=0, + logger=logger) + def __call__(self, img): starttime = time.time() + if self.args.benchmark: + self.autolog.times.start() + ori_im = img.copy() data = {'image': img} data = transform(data, self.preprocess_op) @@ -99,6 +123,8 @@ class TableStructurer(object): return None, 0 img = np.expand_dims(img, axis=0) img = img.copy() + if self.args.benchmark: + self.autolog.times.stamp() if self.use_onnx: input_dict = {} input_dict[self.input_tensor.name] = img @@ -110,6 +136,8 @@ class TableStructurer(object): for output_tensor in self.output_tensors: output = output_tensor.copy_to_cpu() outputs.append(output) + if self.args.benchmark: + self.autolog.times.stamp() preds = {} preds['structure_probs'] = outputs[1] @@ -125,6 +153,8 @@ class TableStructurer(object): '', '', '' ] + structure_str_list + ['
', '', ''] elapse = time.time() - starttime + if self.args.benchmark: + self.autolog.times.end(stamp=True) return (structure_str_list, bbox_list), elapse @@ -164,6 +194,8 @@ def main(args): total_time += elapse count += 1 logger.info("Predict time of {}: {}".format(image_file, elapse)) + if args.benchmark: + table_structurer.autolog.report() if __name__ == "__main__": diff --git a/ppstructure/table/predict_table.py b/ppstructure/table/predict_table.py index fdf611b7ffc049fc745b86233ee127337eaf5f90..8f9c7174904ab3818f62544aeadc97c410070b07 100644 --- a/ppstructure/table/predict_table.py +++ b/ppstructure/table/predict_table.py @@ -14,7 +14,6 @@ import os import sys -import subprocess __dir__ = os.path.dirname(os.path.abspath(__file__)) sys.path.append(__dir__) @@ -61,57 +60,31 @@ class TableSystem(object): self.args = args if not args.show_log: logger.setLevel(logging.INFO) - - self.text_detector = predict_det.TextDetector( - args) if text_detector is None else text_detector - self.text_recognizer = predict_rec.TextRecognizer( - args) if text_recognizer is None else text_recognizer - + args.benchmark = False + self.text_detector = predict_det.TextDetector(copy.deepcopy( + args)) if text_detector is None else text_detector + self.text_recognizer = predict_rec.TextRecognizer(copy.deepcopy( + args)) if text_recognizer is None else text_recognizer + args.benchmark = True self.table_structurer = predict_strture.TableStructurer(args) if args.table_algorithm in ['TableMaster']: self.match = TableMasterMatcher() else: self.match = TableMatch(filter_ocr_result=True) - self.benchmark = args.benchmark self.predictor, self.input_tensor, self.output_tensors, self.config = utility.create_predictor( args, 'table', logger) - if args.benchmark: - import auto_log - pid = os.getpid() - gpu_id = utility.get_infer_gpuid() - self.autolog = auto_log.AutoLogger( - model_name="table", - model_precision=args.precision, - batch_size=1, - data_shape="dynamic", - save_path=None, #args.save_log_path, - inference_config=self.config, - pids=pid, - process_name=None, - gpu_ids=gpu_id if args.use_gpu else None, - time_keys=[ - 'preprocess_time', 'inference_time', 'postprocess_time' - ], - warmup=0, - logger=logger) def __call__(self, img, return_ocr_result_in_table=False): result = dict() time_dict = {'det': 0, 'rec': 0, 'table': 0, 'all': 0, 'match': 0} start = time.time() - if self.args.benchmark: - self.autolog.times.start() structure_res, elapse = self._structure(copy.deepcopy(img)) - if self.benchmark: - self.autolog.times.stamp() result['cell_bbox'] = structure_res[1].tolist() time_dict['table'] = elapse dt_boxes, rec_res, det_elapse, rec_elapse = self._ocr( copy.deepcopy(img)) - if self.benchmark: - self.autolog.times.stamp() time_dict['det'] = det_elapse time_dict['rec'] = rec_elapse @@ -126,8 +99,6 @@ class TableSystem(object): result['html'] = pred_html end = time.time() time_dict['all'] = end - start - if self.benchmark: - self.autolog.times.end(stamp=True) return result, time_dict def _structure(self, img): @@ -233,12 +204,13 @@ def main(args): f_html.close() if args.benchmark: - table_sys.autolog.report() + table_sys.table_structurer.autolog.report() if __name__ == "__main__": args = parse_args() if args.use_mp: + import subprocess p_list = [] total_process_num = args.total_process_num for process_id in range(total_process_num):