diff --git a/configs/rec/ch_ppocr_v2.0/rec_chinese_common_train_v2.0.yml b/configs/rec/ch_ppocr_v2.0/rec_chinese_common_train_v2.0.yml index 6a524e22cf4dea4c573d6b67e752c8527e973185..717c16814bac2f6fca78aa63566df12bd8cbf67b 100644 --- a/configs/rec/ch_ppocr_v2.0/rec_chinese_common_train_v2.0.yml +++ b/configs/rec/ch_ppocr_v2.0/rec_chinese_common_train_v2.0.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: True + save_res_path: ./output/rec/predicts_chinese_common_v2.0.txt Optimizer: diff --git a/configs/rec/ch_ppocr_v2.0/rec_chinese_lite_train_v2.0.yml b/configs/rec/ch_ppocr_v2.0/rec_chinese_lite_train_v2.0.yml index c96621c5684f3861a7bc2f5aa8d9684e6512e228..d60832b7c9151a60d580e9ccb3b264cc685faf0f 100644 --- a/configs/rec/ch_ppocr_v2.0/rec_chinese_lite_train_v2.0.yml +++ b/configs/rec/ch_ppocr_v2.0/rec_chinese_lite_train_v2.0.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: True + # save_res_path: ./output/rec/predicts_chinese_lite_v2.0.txt Optimizer: diff --git a/configs/rec/rec_icdar15_train.yml b/configs/rec/rec_icdar15_train.yml index 5ae47c67d8b062746d422daac44011fb5aca38e2..79e3ff88869d9e0fc7be5563830eed1244e38b76 100644 --- a/configs/rec/rec_icdar15_train.yml +++ b/configs/rec/rec_icdar15_train.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_ic15.txt Optimizer: name: Adam diff --git a/configs/rec/rec_mv3_none_bilstm_ctc.yml b/configs/rec/rec_mv3_none_bilstm_ctc.yml index 900e98b6b34de824d9afb91b1867a86fe2debc24..9e0bd23edba053b44fc7241c0a587ced5cd1ac76 100644 --- a/configs/rec/rec_mv3_none_bilstm_ctc.yml +++ b/configs/rec/rec_mv3_none_bilstm_ctc.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_mv3_none_bilstm_ctc.txt Optimizer: name: Adam diff --git a/configs/rec/rec_mv3_none_none_ctc.yml b/configs/rec/rec_mv3_none_none_ctc.yml index 6d86b90c007cab9708cab6db6d8e3045dd5187fb..904afe1134b565d6459cdcda4cbfa43ae4925b92 100644 --- a/configs/rec/rec_mv3_none_none_ctc.yml +++ b/configs/rec/rec_mv3_none_none_ctc.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_mv3_none_none_ctc.txt Optimizer: name: Adam diff --git a/configs/rec/rec_mv3_tps_bilstm_att.yml b/configs/rec/rec_mv3_tps_bilstm_att.yml index 33aed74d83f9ab18b7bdfb12d0ed315c6eebc010..feaeb0545c687774938521e4c45c026207172f11 100644 --- a/configs/rec/rec_mv3_tps_bilstm_att.yml +++ b/configs/rec/rec_mv3_tps_bilstm_att.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_mv3_tps_bilstm_att.txt Optimizer: diff --git a/configs/rec/rec_mv3_tps_bilstm_ctc.yml b/configs/rec/rec_mv3_tps_bilstm_ctc.yml index 026c6a9dfbd6b6b543c0b4260c43cbf98e192e7b..65ab23c42aff54ee548867e3482d7400603551ad 100644 --- a/configs/rec/rec_mv3_tps_bilstm_ctc.yml +++ b/configs/rec/rec_mv3_tps_bilstm_ctc.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_mv3_tps_bilstm_ctc.txt Optimizer: name: Adam diff --git a/configs/rec/rec_r34_vd_none_bilstm_ctc.yml b/configs/rec/rec_r34_vd_none_bilstm_ctc.yml index 4052d426e51aa8c6e82ec216cfd65226922be602..331bb36ed84b83dc62a0f9b15524457238dedc13 100644 --- a/configs/rec/rec_r34_vd_none_bilstm_ctc.yml +++ b/configs/rec/rec_r34_vd_none_bilstm_ctc.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_r34_vd_none_bilstm_ctc.txt Optimizer: name: Adam diff --git a/configs/rec/rec_r34_vd_none_none_ctc.yml b/configs/rec/rec_r34_vd_none_none_ctc.yml index c3e1d9a3a91ab6a51e28d458623aea788b952ca0..695a46958f669e4cb9508646080b45ac0767b8c9 100644 --- a/configs/rec/rec_r34_vd_none_none_ctc.yml +++ b/configs/rec/rec_r34_vd_none_none_ctc.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_r34_vd_none_none_ctc.txt Optimizer: name: Adam diff --git a/configs/rec/rec_r34_vd_tps_bilstm_att.yml b/configs/rec/rec_r34_vd_tps_bilstm_att.yml index 87a14559849abd0a47c45f52b80d06bb0790ef0e..fdd3588c844ffd7ed61de73077ae2994f0ad498d 100644 --- a/configs/rec/rec_r34_vd_tps_bilstm_att.yml +++ b/configs/rec/rec_r34_vd_tps_bilstm_att.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_b3_rare_r34_none_gru.txt Optimizer: diff --git a/configs/rec/rec_r34_vd_tps_bilstm_ctc.yml b/configs/rec/rec_r34_vd_tps_bilstm_ctc.yml index 9c51962e58763838eb4a101d1d52a1a030e4b643..67108a6eaca2dd6f239261f5184341e5ade00dc0 100644 --- a/configs/rec/rec_r34_vd_tps_bilstm_ctc.yml +++ b/configs/rec/rec_r34_vd_tps_bilstm_ctc.yml @@ -19,6 +19,7 @@ Global: max_text_length: 25 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_r34_vd_tps_bilstm_ctc.txt Optimizer: name: Adam diff --git a/configs/rec/rec_r50_fpn_srn.yml b/configs/rec/rec_r50_fpn_srn.yml index 34a997f3c3f6bc3dccdc62757c6adb1e1d17cc53..fa7b1ae4e5fed41d3aa3670d6672cca01b63c359 100644 --- a/configs/rec/rec_r50_fpn_srn.yml +++ b/configs/rec/rec_r50_fpn_srn.yml @@ -20,6 +20,7 @@ Global: num_heads: 8 infer_mode: False use_space_char: False + save_res_path: ./output/rec/predicts_srn.txt Optimizer: diff --git a/tools/infer_rec.py b/tools/infer_rec.py index 075ec261e492cf21c668364ae6119fb4903f823b..2563f5a8197ed39b1b5d44c7cfee32797e760758 100755 --- a/tools/infer_rec.py +++ b/tools/infer_rec.py @@ -73,35 +73,45 @@ def main(): global_config['infer_mode'] = True ops = create_operators(transforms, global_config) + save_res_path = config['Global'].get('save_res_path', + "./output/rec/predicts_rec.txt") + if not os.path.exists(os.path.dirname(save_res_path)): + os.makedirs(os.path.dirname(save_res_path)) + model.eval() - for file in get_image_file_list(config['Global']['infer_img']): - logger.info("infer_img: {}".format(file)) - with open(file, 'rb') as f: - img = f.read() - data = {'image': img} - batch = transform(data, ops) - if config['Architecture']['algorithm'] == "SRN": - encoder_word_pos_list = np.expand_dims(batch[1], axis=0) - gsrm_word_pos_list = np.expand_dims(batch[2], axis=0) - gsrm_slf_attn_bias1_list = np.expand_dims(batch[3], axis=0) - gsrm_slf_attn_bias2_list = np.expand_dims(batch[4], axis=0) - - others = [ - paddle.to_tensor(encoder_word_pos_list), - paddle.to_tensor(gsrm_word_pos_list), - paddle.to_tensor(gsrm_slf_attn_bias1_list), - paddle.to_tensor(gsrm_slf_attn_bias2_list) - ] - - images = np.expand_dims(batch[0], axis=0) - images = paddle.to_tensor(images) - if config['Architecture']['algorithm'] == "SRN": - preds = model(images, others) - else: - preds = model(images) - post_result = post_process_class(preds) - for rec_reuslt in post_result: - logger.info('\t result: {}'.format(rec_reuslt)) + + with open(save_res_path, "w") as fout: + for file in get_image_file_list(config['Global']['infer_img']): + logger.info("infer_img: {}".format(file)) + with open(file, 'rb') as f: + img = f.read() + data = {'image': img} + batch = transform(data, ops) + if config['Architecture']['algorithm'] == "SRN": + encoder_word_pos_list = np.expand_dims(batch[1], axis=0) + gsrm_word_pos_list = np.expand_dims(batch[2], axis=0) + gsrm_slf_attn_bias1_list = np.expand_dims(batch[3], axis=0) + gsrm_slf_attn_bias2_list = np.expand_dims(batch[4], axis=0) + + others = [ + paddle.to_tensor(encoder_word_pos_list), + paddle.to_tensor(gsrm_word_pos_list), + paddle.to_tensor(gsrm_slf_attn_bias1_list), + paddle.to_tensor(gsrm_slf_attn_bias2_list) + ] + + images = np.expand_dims(batch[0], axis=0) + images = paddle.to_tensor(images) + if config['Architecture']['algorithm'] == "SRN": + preds = model(images, others) + else: + preds = model(images) + post_result = post_process_class(preds) + for rec_reuslt in post_result: + logger.info('\t result: {}'.format(rec_reuslt)) + if len(rec_reuslt) >= 2: + fout.write(file + "\t" + rec_reuslt[0] + "\t" + str( + rec_reuslt[1]) + "\n") logger.info("success!")