modify batch num in rec and fix sorted_box func

This commit is contained in:
LDOUBLEV 2020-05-20 16:19:49 +08:00
parent 0bf34a2577
commit 2bd6323642
2 changed files with 5 additions and 2 deletions

View File

@ -31,6 +31,7 @@ class TextRecognizer(object):
image_shape = [int(v) for v in args.rec_image_shape.split(",")] image_shape = [int(v) for v in args.rec_image_shape.split(",")]
self.rec_image_shape = image_shape self.rec_image_shape = image_shape
self.character_type = args.rec_char_type self.character_type = args.rec_char_type
self.rec_batch_num = args.rec_batch_num
char_ops_params = {} char_ops_params = {}
char_ops_params["character_type"] = args.rec_char_type char_ops_params["character_type"] = args.rec_char_type
char_ops_params["character_dict_path"] = args.rec_char_dict_path char_ops_params["character_dict_path"] = args.rec_char_dict_path
@ -59,8 +60,8 @@ class TextRecognizer(object):
def __call__(self, img_list): def __call__(self, img_list):
img_num = len(img_list) img_num = len(img_list)
batch_num = 12
rec_res = [] rec_res = []
batch_num = self.rec_batch_num
predict_time = 0 predict_time = 0
for beg_img_no in range(0, img_num, batch_num): for beg_img_no in range(0, img_num, batch_num):
end_img_no = min(img_num, beg_img_no + batch_num) end_img_no = min(img_num, beg_img_no + batch_num)

View File

@ -56,6 +56,7 @@ def parse_args():
parser.add_argument("--rec_model_dir", type=str) parser.add_argument("--rec_model_dir", type=str)
parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320") 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_char_type", type=str, default='ch')
parser.add_argument("--rec_batch_num", type=int, default=30)
parser.add_argument( parser.add_argument(
"--rec_char_dict_path", "--rec_char_dict_path",
type=str, type=str,
@ -172,7 +173,8 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
continue continue
font = ImageFont.truetype( font = ImageFont.truetype(
"./doc/simfang.ttf", font_size, encoding="utf-8") "./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( draw_txt.text(
(20, gap * (count + 1)), new_txt, txt_color, font=font) (20, gap * (count + 1)), new_txt, txt_color, font=font)
count += 1 count += 1