diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index ebf0b0ba2aaf33c05307f7c2d70bce1c7124e15f..e0f2c41fa2aba23491efee920afbd76db1ec84e0 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -34,7 +34,7 @@ import tools.infer.predict_det as predict_det import tools.infer.predict_cls as predict_cls from ppocr.utils.utility import get_image_file_list, check_and_read from ppocr.utils.logging import get_logger -from tools.infer.utility import draw_ocr_box_txt2, get_rotate_crop_image +from tools.infer.utility import draw_ocr_box_txt, get_rotate_crop_image logger = get_logger() @@ -189,7 +189,7 @@ 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_box_txt2( + draw_img = draw_ocr_box_txt( image, boxes, txts, diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 8045ec44810175ae236f1c44a86cd85cf09bbd8b..8ae8822f760b6e5fb7947ac878ce798ebe22883d 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -397,62 +397,10 @@ def draw_ocr(image, def draw_ocr_box_txt(image, boxes, - txts, + txts=None, 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)) - - import random - - random.seed(0) - draw_left = ImageDraw.Draw(img_left) - draw_right = ImageDraw.Draw(img_right) - 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) - draw_right.polygon( - [ - box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], - box[2][1], box[3][0], box[3][1] - ], - outline=color) - box_height = math.sqrt((box[0][0] - box[3][0])**2 + (box[0][1] - box[3][ - 1])**2) - box_width = math.sqrt((box[0][0] - box[1][0])**2 + (box[0][1] - box[1][ - 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") - cur_y = box[0][1] - for c in txt: - char_size = font.getsize(c) - draw_right.text( - (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font) - 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") - 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) - img_show = Image.new('RGB', (w * 2, h), (255, 255, 255)) - img_show.paste(img_left, (0, 0, w, h)) - img_show.paste(img_right, (w, 0, w * 2, h)) - return np.array(img_show) - - -def draw_ocr_box_txt2(image, - boxes, - txts=None, - scores=None, - drop_score=0.5, - font_path="./doc/fonts/simfang.ttf"): + font_path="./doc/fonts/simfang.ttf"): h, w = image.height, image.width img_left = image.copy() img_right = np.ones((h, w, 3), dtype=np.uint8) * 255