From aa7e9ac34e219394b39531835f76e2b8f2986740 Mon Sep 17 00:00:00 2001 From: tink2123 Date: Thu, 3 Sep 2020 16:06:57 +0800 Subject: [PATCH] polish code --- tools/infer/predict_rec.py | 56 ++++++++++++++++++++------------------ tools/infer/utility.py | 10 +++---- tools/infer_rec.py | 5 +--- 3 files changed, 35 insertions(+), 36 deletions(-) diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index 8f676294..e592ccdc 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -196,32 +196,36 @@ class TextRecognizer(object): norm_img_batch.append(norm_img[0]) norm_img_batch = np.concatenate(norm_img_batch, axis=0) - - encoder_word_pos_list = np.concatenate(encoder_word_pos_list) - - gsrm_word_pos_list = np.concatenate(gsrm_word_pos_list) - - gsrm_slf_attn_bias1_list = np.concatenate(gsrm_slf_attn_bias1_list) - - gsrm_slf_attn_bias2_list = np.concatenate(gsrm_slf_attn_bias2_list) - - starttime = time.time() - - norm_img_batch = fluid.core.PaddleTensor(norm_img_batch) - encoder_word_pos_list = fluid.core.PaddleTensor( - encoder_word_pos_list) - gsrm_word_pos_list = fluid.core.PaddleTensor(gsrm_word_pos_list) - gsrm_slf_attn_bias1_list = fluid.core.PaddleTensor( - gsrm_slf_attn_bias1_list) - gsrm_slf_attn_bias2_list = fluid.core.PaddleTensor( - gsrm_slf_attn_bias2_list) - - inputs = [ - norm_img_batch, encoder_word_pos_list, gsrm_slf_attn_bias1_list, - gsrm_slf_attn_bias2_list, gsrm_word_pos_list - ] - - self.predictor.run(inputs) + norm_img_batch = norm_img_batch.copy() + + if self.loss_type == "srn": + encoder_word_pos_list = np.concatenate(encoder_word_pos_list) + gsrm_word_pos_list = np.concatenate(gsrm_word_pos_list) + gsrm_slf_attn_bias1_list = np.concatenate( + gsrm_slf_attn_bias1_list) + gsrm_slf_attn_bias2_list = np.concatenate( + gsrm_slf_attn_bias2_list) + starttime = time.time() + + norm_img_batch = fluid.core.PaddleTensor(norm_img_batch) + encoder_word_pos_list = fluid.core.PaddleTensor( + encoder_word_pos_list) + gsrm_word_pos_list = fluid.core.PaddleTensor(gsrm_word_pos_list) + gsrm_slf_attn_bias1_list = fluid.core.PaddleTensor( + gsrm_slf_attn_bias1_list) + gsrm_slf_attn_bias2_list = fluid.core.PaddleTensor( + gsrm_slf_attn_bias2_list) + + inputs = [ + norm_img_batch, encoder_word_pos_list, + gsrm_slf_attn_bias1_list, gsrm_slf_attn_bias2_list, + gsrm_word_pos_list + ] + + self.predictor.run(inputs) + else: + self.input_tensor.copy_from_cpu(norm_img_batch) + self.predictor.zero_copy_run() if self.loss_type == "ctc": rec_idx_batch = self.output_tensors[0].copy_to_cpu() diff --git a/tools/infer/utility.py b/tools/infer/utility.py index e4e2a9fc..ccf6e32b 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -59,10 +59,10 @@ def parse_args(): parser.add_argument("--det_sast_polygon", type=bool, default=False) #params for text recognizer - parser.add_argument("--rec_algorithm", type=str, default='SRN') + parser.add_argument("--rec_algorithm", type=str, default='CRNN') parser.add_argument("--rec_model_dir", type=str) - parser.add_argument("--rec_image_shape", type=str, default="3, 64, 256") - parser.add_argument("--rec_char_type", type=str, default='en') + parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320") + parser.add_argument("--rec_char_type", type=str, default='ch') parser.add_argument("--rec_batch_num", type=int, default=30) parser.add_argument("--max_text_length", type=int, default=25) parser.add_argument( @@ -107,11 +107,9 @@ def create_predictor(args, mode): # use zero copy config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass") - #config.switch_use_feed_fetch_ops(False) - config.switch_use_feed_fetch_ops(True) + config.switch_use_feed_fetch_ops(False) predictor = create_paddle_predictor(config) input_names = predictor.get_input_names() - print(input_names) for name in input_names: input_tensor = predictor.get_input_tensor(name) output_names = predictor.get_output_names() diff --git a/tools/infer_rec.py b/tools/infer_rec.py index e4b75e86..29fc5b40 100755 --- a/tools/infer_rec.py +++ b/tools/infer_rec.py @@ -162,10 +162,7 @@ def main(): fluid.io.save_inference_model( "./output/", - feeded_var_names=[ - 'image', 'encoder_word_pos', 'gsrm_slf_attn_bias1', - 'gsrm_slf_attn_bias2', 'gsrm_word_pos' - ], + feeded_var_names=['image'], target_vars=target_var, executor=exe, main_program=eval_prog, -- GitLab