fix visualize

This commit is contained in:
littletomatodonkey 2020-09-22 07:07:42 +00:00
parent 250412a408
commit 571275e03c
2 changed files with 19 additions and 11 deletions

View File

@ -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)

View File

@ -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)