From f6698a32d8fbeaa4fb00bba5c5e9f652b02e1c1b Mon Sep 17 00:00:00 2001
From: WenmuZhou <572459439@qq.com>
Date: Mon, 5 Sep 2022 15:02:22 +0800
Subject: [PATCH] support onnx infer of SLANet

---
 ppocr/modeling/heads/table_att_head.py | 16 +++++++---------
 ppstructure/table/predict_structure.py | 19 ++++++++++++-------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/ppocr/modeling/heads/table_att_head.py b/ppocr/modeling/heads/table_att_head.py
index 00b43410..2e077124 100644
--- a/ppocr/modeling/heads/table_att_head.py
+++ b/ppocr/modeling/heads/table_att_head.py
@@ -166,6 +166,7 @@ class SLAHead(nn.Layer):
         self.max_text_length = max_text_length
         self.emb = self._char_to_onehot
         self.num_embeddings = out_channels
+        self.loc_reg_num = loc_reg_num
 
         # structure
         self.structure_attention_cell = AttentionGRUCell(
@@ -213,16 +214,15 @@ class SLAHead(nn.Layer):
         fea = fea.transpose([0, 2, 1])  # (NTC)(batch, width, channels)
 
         hidden = paddle.zeros((batch_size, self.hidden_size))
-        structure_preds = []
-        loc_preds = []
+        structure_preds = paddle.zeros((batch_size, self.max_text_length + 1, self.num_embeddings))
+        loc_preds = paddle.zeros((batch_size, self.max_text_length + 1, self.loc_reg_num))
         if self.training and targets is not None:
             structure = targets[0]
             for i in range(self.max_text_length + 1):
                 hidden, structure_step, loc_step = self._decode(structure[:, i],
                                                                 fea, hidden)
-                structure_preds.append(structure_step)
-                loc_preds.append(loc_step)
-        else:
+                structure_preds[:, i, :] = structure_step
+                loc_preds[:, i, :] = loc_step
             pre_chars = paddle.zeros(shape=[batch_size], dtype="int32")
             max_text_length = paddle.to_tensor(self.max_text_length)
             # for export
@@ -231,10 +231,8 @@ class SLAHead(nn.Layer):
                 hidden, structure_step, loc_step = self._decode(pre_chars, fea,
                                                                 hidden)
                 pre_chars = structure_step.argmax(axis=1, dtype="int32")
-                structure_preds.append(structure_step)
-                loc_preds.append(loc_step)
-        structure_preds = paddle.stack(structure_preds, axis=1)
-        loc_preds = paddle.stack(loc_preds, axis=1)
+                structure_preds[:, i, :] = structure_step
+                loc_preds[:, i, :] = loc_step
         if not self.training:
             structure_preds = F.softmax(structure_preds)
         return {'structure_probs': structure_preds, 'loc_preds': loc_preds}
diff --git a/ppstructure/table/predict_structure.py b/ppstructure/table/predict_structure.py
index 45cbba3e..0bf10085 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.use_onnx = args.use_onnx
         pre_process_list = build_pre_process_list(args)
         if args.table_algorithm not in ['TableMaster']:
             postprocess_params = {
@@ -98,13 +99,17 @@ class TableStructurer(object):
             return None, 0
         img = np.expand_dims(img, axis=0)
         img = img.copy()
-
-        self.input_tensor.copy_from_cpu(img)
-        self.predictor.run()
-        outputs = []
-        for output_tensor in self.output_tensors:
-            output = output_tensor.copy_to_cpu()
-            outputs.append(output)
+        if self.use_onnx:
+            input_dict = {}
+            input_dict[self.input_tensor.name] = img
+            outputs = self.predictor.run(self.output_tensors, input_dict)
+        else:
+            self.input_tensor.copy_from_cpu(img)
+            self.predictor.run()
+            outputs = []
+            for output_tensor in self.output_tensors:
+                output = output_tensor.copy_to_cpu()
+                outputs.append(output)
 
         preds = {}
         preds['structure_probs'] = outputs[1]
-- 
GitLab