layout recognition refinement onnx support (#12068)

* layout recognition refinement onnx support * fix codestyle
PaddlePaddle · May 9, 2024 · 5818196 · 5818196
1 parent 95e3103
commit 5818196
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 18 deletions.
diff --git a/paddleocr.py b/paddleocr.py
@@ -812,10 +812,11 @@ def __init__(self, **kwargs):
             layout_model_config["url"],
         )
         # download model
-        maybe_download(params.det_model_dir, det_url)
-        maybe_download(params.rec_model_dir, rec_url)
-        maybe_download(params.table_model_dir, table_url)
-        maybe_download(params.layout_model_dir, layout_url)
+        if not params.use_onnx:
+            maybe_download(params.det_model_dir, det_url)
+            maybe_download(params.rec_model_dir, rec_url)
+            maybe_download(params.table_model_dir, table_url)
+            maybe_download(params.layout_model_dir, layout_url)
 
         if params.rec_char_dict_path is None:
             params.rec_char_dict_path = str(

diff --git a/ppstructure/layout/predict_layout.py b/ppstructure/layout/predict_layout.py
@@ -65,6 +65,7 @@ def __init__(self, args):
             self.output_tensors,
             self.config,
         ) = utility.create_predictor(args, "layout", logger)
+        self.use_onnx = args.use_onnx
 
     def __call__(self, img):
         ori_im = img.copy()
@@ -81,21 +82,31 @@ def __call__(self, img):
         preds, elapse = 0, 1
         starttime = time.time()
 
-        self.input_tensor.copy_from_cpu(img)
-        self.predictor.run()
-
         np_score_list, np_boxes_list = [], []
-        output_names = self.predictor.get_output_names()
-        num_outs = int(len(output_names) / 2)
-        for out_idx in range(num_outs):
-            np_score_list.append(
-                self.predictor.get_output_handle(output_names[out_idx]).copy_to_cpu()
-            )
-            np_boxes_list.append(
-                self.predictor.get_output_handle(
-                    output_names[out_idx + num_outs]
-                ).copy_to_cpu()
-            )
+        if self.use_onnx:
+            input_dict = {}
+            input_dict[self.input_tensor.name] = img
+            outputs = self.predictor.run(self.output_tensors, input_dict)
+            num_outs = int(len(outputs) / 2)
+            for out_idx in range(num_outs):
+                np_score_list.append(outputs[out_idx])
+                np_boxes_list.append(outputs[out_idx + num_outs])
+        else:
+            self.input_tensor.copy_from_cpu(img)
+            self.predictor.run()
+            output_names = self.predictor.get_output_names()
+            num_outs = int(len(output_names) / 2)
+            for out_idx in range(num_outs):
+                np_score_list.append(
+                    self.predictor.get_output_handle(
+                        output_names[out_idx]
+                    ).copy_to_cpu()
+                )
+                np_boxes_list.append(
+                    self.predictor.get_output_handle(
+                        output_names[out_idx + num_outs]
+                    ).copy_to_cpu()
+                )
         preds = dict(boxes=np_score_list, boxes_num=np_boxes_list)
 
         post_preds = self.postprocess_op(ori_im, img, preds)