diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index a3f68f04..6fa51e70 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -31,6 +31,7 @@ class TextRecognizer(object): image_shape = [int(v) for v in args.rec_image_shape.split(",")] self.rec_image_shape = image_shape self.character_type = args.rec_char_type + self.rec_batch_num = args.rec_batch_num char_ops_params = {} char_ops_params["character_type"] = args.rec_char_type char_ops_params["character_dict_path"] = args.rec_char_dict_path @@ -59,8 +60,8 @@ class TextRecognizer(object): def __call__(self, img_list): img_num = len(img_list) - batch_num = 12 rec_res = [] + batch_num = self.rec_batch_num predict_time = 0 for beg_img_no in range(0, img_num, batch_num): end_img_no = min(img_num, beg_img_no + batch_num) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index fd72a19f..947a5495 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -56,6 +56,7 @@ def parse_args(): parser.add_argument("--rec_model_dir", type=str) 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( "--rec_char_dict_path", type=str, @@ -172,7 +173,8 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5): continue font = ImageFont.truetype( "./doc/simfang.ttf", font_size, encoding="utf-8") - new_txt = str(count) + ': ' + txt + ' ' + '%.3f' % (scores[count]) + new_txt = str(count) + ': ' + txt + ' ' + '%.3f' % ( + scores[count]) draw_txt.text( (20, gap * (count + 1)), new_txt, txt_color, font=font) count += 1