diff --git a/configs/table/SLANet.yml b/configs/table/SLANet.yml index 384c95852e815f9780328f63cbbd52fa0ef3deb4..a896614556e36f77bd784218b6c2f29914219dbe 100644 --- a/configs/table/SLANet.yml +++ b/configs/table/SLANet.yml @@ -12,7 +12,7 @@ Global: checkpoints: save_inference_dir: ./output/SLANet/infer use_visualdl: False - infer_img: doc/table/table.jpg + infer_img: ppstructure/docs/table/table.jpg # for data or label process character_dict_path: ppocr/utils/dict/table_structure_dict.txt character_type: en diff --git a/configs/table/SLANet_ch.yml b/configs/table/SLANet_ch.yml index a3fc1c68ddd2287e6bbaa6c53c5e20961950d23f..3b1e5c6bd9dd4cd2a084d557a1285983a56bdf2a 100644 --- a/configs/table/SLANet_ch.yml +++ b/configs/table/SLANet_ch.yml @@ -12,7 +12,7 @@ Global: checkpoints: save_inference_dir: ./output/SLANet_ch/infer use_visualdl: False - infer_img: doc/table/table.jpg + infer_img: ppstructure/docs/table/table.jpg # for data or label process character_dict_path: ppocr/utils/dict/table_structure_dict_ch.txt character_type: en diff --git a/ppocr/modeling/heads/table_att_head.py b/ppocr/modeling/heads/table_att_head.py index 00b434105bd9fe1f0d928c5f026dc5804b33fe23..d3c86e22b02e08c18d8d5cb193f2ffb8b07ad785 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,15 +214,17 @@ 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)) + structure_preds.stop_gradient = True + loc_preds.stop_gradient = True 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) + structure_preds[:, i, :] = structure_step + loc_preds[:, i, :] = loc_step else: pre_chars = paddle.zeros(shape=[batch_size], dtype="int32") max_text_length = paddle.to_tensor(self.max_text_length) @@ -231,10 +234,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]