diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index 29c4d7e8e35ceda3966dfcadcca5f0ae985d1bb1..766cc577b07bc9845af610a7a7e1ba50614d64eb 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -160,8 +160,13 @@ def main(args): txts = [rec_res[i][0] for i in range(len(rec_res))] scores = [rec_res[i][1] for i in range(len(rec_res))] - draw_img = draw_ocr( - image, boxes, txts, scores, drop_score=drop_score, font_path=font_path) + draw_img = draw_ocr_box_txt( + image, + boxes, + txts, + scores, + drop_score=drop_score, + font_path=font_path) draw_img_save = "./inference_results/" if not os.path.exists(draw_img_save): os.makedirs(draw_img_save) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 45d7b73707904d3bb2df665a1cf348a32c70f852..1eb57d34490edd3557dbe09a541b0a173ac7bbdf 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -72,9 +72,7 @@ def parse_args(): default="./ppocr/utils/ppocr_keys_v1.txt") parser.add_argument("--use_space_char", type=str2bool, default=True) parser.add_argument( - "--vis_font_path", - type=str, - default="./doc/simfang.ttf") + "--vis_font_path", type=str, default="./doc/simfang.ttf") # params for text classifier parser.add_argument("--use_angle_cls", type=str2bool, default=False) @@ -203,7 +201,12 @@ def draw_ocr(image, return image -def draw_ocr_box_txt(image, boxes, txts, font_path="./doc/simfang.ttf"): +def draw_ocr_box_txt(image, + boxes, + txts, + scores=None, + drop_score=0.5, + font_path="./doc/simfang.ttf"): h, w = image.height, image.width img_left = image.copy() img_right = Image.new('RGB', (w, h), (255, 255, 255)) @@ -213,7 +216,9 @@ def draw_ocr_box_txt(image, boxes, txts, font_path="./doc/simfang.ttf"): random.seed(0) draw_left = ImageDraw.Draw(img_left) draw_right = ImageDraw.Draw(img_right) - for (box, txt) in zip(boxes, txts): + for idx, (box, txt) in enumerate(zip(boxes, txts)): + if scores is not None and scores[idx] < drop_score: + continue color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) draw_left.polygon(box, fill=color) @@ -229,8 +234,7 @@ def draw_ocr_box_txt(image, boxes, txts, font_path="./doc/simfang.ttf"): 1])**2) if box_height > 2 * box_width: font_size = max(int(box_width * 0.9), 10) - font = ImageFont.truetype( - font_path, font_size, encoding="utf-8") + font = ImageFont.truetype(font_path, font_size, encoding="utf-8") cur_y = box[0][1] for c in txt: char_size = font.getsize(c) @@ -239,8 +243,7 @@ def draw_ocr_box_txt(image, boxes, txts, font_path="./doc/simfang.ttf"): cur_y += char_size[1] else: font_size = max(int(box_height * 0.8), 10) - font = ImageFont.truetype( - font_path, font_size, encoding="utf-8") + font = ImageFont.truetype(font_path, font_size, encoding="utf-8") draw_right.text( [box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font) img_left = Image.blend(image, img_left, 0.5)