diff --git a/ppocr/modeling/heads/table_att_head.py b/ppocr/modeling/heads/table_att_head.py index 00b434105bd9fe1f0d928c5f026dc5804b33fe23..2e0771243ddba4d6d453139ab7c71e96e4f56aae 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 45cbba3e298004d3711b05e6fb7cffecae637601..0bf100852b9e9d501dfc858d8ce0787da42a61ed 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]