fix table recogition benckmark error

4cf04cbe · 文幕地方 · 4078b0fe · 4cf04cbe · 4cf04cbe
隐藏空白更改
内联并排

Showing with 40 addition and 36 deletion

ppstructure/table/predict_structure.py ppstructure/table/predict_structure.py +32 -0

ppstructure/table/predict_table.py ppstructure/table/predict_table.py +8 -36

未找到文件。
--- 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):
            '<html>', '<body>', '<table>'
        ] + structure_str_list + ['</table>', '</body>', '</html>']
        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__":

--- 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):