diff --git a/ppstructure/table/predict_table.py b/ppstructure/table/predict_table.py index 402d6c24189d044e2ee6d359edef8624d4aae145..d7d1760d3567760d3d230769b3fb42a379109652 100644 --- a/ppstructure/table/predict_table.py +++ b/ppstructure/table/predict_table.py @@ -28,6 +28,7 @@ import numpy as np import time import tools.infer.predict_rec as predict_rec import tools.infer.predict_det as predict_det +import tools.infer.utility as utility from ppocr.utils.utility import get_image_file_list, check_and_read_gif from ppocr.utils.logging import get_logger from ppstructure.table.matcher import distance, compute_iou @@ -59,11 +60,37 @@ class TableSystem(object): self.text_recognizer = predict_rec.TextRecognizer( args) if text_recognizer is None else text_recognizer self.table_structurer = predict_strture.TableStructurer(args) + 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() ori_im = img.copy() + if self.benchmark: + self.autolog.times.start() structure_res, elapse = self.table_structurer(copy.deepcopy(img)) + if self.benchmark: + self.autolog.times.stamp() dt_boxes, elapse = self.text_detector(copy.deepcopy(img)) dt_boxes = sorted_boxes(dt_boxes) if return_ocr_result_in_table: @@ -77,13 +104,11 @@ class TableSystem(object): box = [x_min, y_min, x_max, y_max] r_boxes.append(box) dt_boxes = np.array(r_boxes) - logger.debug("dt_boxes num : {}, elapse : {}".format( len(dt_boxes), elapse)) if dt_boxes is None: return None, None img_crop_list = [] - for i in range(len(dt_boxes)): det_box = dt_boxes[i] x0, y0, x1, y1 = expand(2, det_box, ori_im.shape) @@ -92,10 +117,14 @@ class TableSystem(object): rec_res, elapse = self.text_recognizer(img_crop_list) logger.debug("rec_res num : {}, elapse : {}".format( len(rec_res), elapse)) + if self.benchmark: + self.autolog.times.stamp() if return_ocr_result_in_table: result['rec_res'] = rec_res pred_html, pred = self.rebuild_table(structure_res, dt_boxes, rec_res) result['html'] = pred_html + if self.benchmark: + self.autolog.times.end(stamp=True) return result def rebuild_table(self, structure_res, dt_boxes, rec_res): @@ -213,11 +242,15 @@ def main(args): logger.info('excel saved to {}'.format(excel_path)) elapse = time.time() - starttime logger.info("Predict time : {:.3f}s".format(elapse)) + if args.benchmark: + text_sys.autolog.report() + print('ok') if __name__ == "__main__": args = parse_args() if args.use_mp: + print('mp') p_list = [] total_process_num = args.total_process_num for process_id in range(total_process_num):