fix inference KIE model using onnx model

Author: Alex37882388

ppstructure/kie/predict_kie_token_ser.py

@@ -40,6 +40,7 @@
 
 class SerPredictor(object):
     def __init__(self, args):
+        self.args = args
         self.ocr_engine = PaddleOCR(
             use_angle_cls=args.use_angle_cls,
             det_model_dir=args.det_model_dir,
@@ -113,15 +114,18 @@ def __call__(self, img):
                 data[idx] = np.expand_dims(data[idx], axis=0)
             else:
                 data[idx] = [data[idx]]
+        if self.args.use_onnx:
+            input_tensor = {name: data[idx] for idx, name in enumerate(self.input_tensor)}
+            self.output_tensors = self.predictor.run(None, input_tensor)
+        else:
+            for idx in range(len(self.input_tensor)):
+                self.input_tensor[idx].copy_from_cpu(data[idx])
 
-        for idx in range(len(self.input_tensor)):
-            self.input_tensor[idx].copy_from_cpu(data[idx])
-
-        self.predictor.run()
+            self.predictor.run()
 
         outputs = []
         for output_tensor in self.output_tensors:
-            output = output_tensor.copy_to_cpu()
+            output = output_tensor if self.args.use_onnx else output_tensor.copy_to_cpu()
             outputs.append(output)
         preds = outputs[0]
 

tools/infer/utility.py

@@ -221,7 +221,7 @@ def create_predictor(args, mode, logger):
                 providers=["CPUExecutionProvider"],
                 sess_options=sess_options,
             )
-        return sess, sess.get_inputs()[0], None, None
+        return sess, [vo.name for vo in sess.get_inputs()], None, None
 
     else:
         file_names = ["model", "inference"]