提交 e0b641cc 编写于 作者: L LDOUBLEV

fix visual bug

上级 00e9e079
...@@ -140,22 +140,21 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5): ...@@ -140,22 +140,21 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
return(array): return(array):
the visualized img the visualized img
""" """
img = image
if scores is None: if scores is None:
scores = [1] * len(boxes) scores = [1] * len(boxes)
for (box, score) in zip(boxes, scores): for (box, score) in zip(boxes, scores):
if score < drop_score or math.isnan(score): if score < drop_score or math.isnan(score):
continue continue
box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64) box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64)
img = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 3) image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
if draw_txt: if draw_txt:
img = np.array(resize_img(img, input_size=600)) img = np.array(resize_img(image, input_size=600))
txt_img = text_visual( txt_img = text_visual(
txts, scores, img_h=img.shape[0], img_w=600, threshold=drop_score) txts, scores, img_h=img.shape[0], img_w=600, threshold=drop_score)
img = np.concatenate([np.array(img), np.array(txt_img)], axis=1) img = np.concatenate([np.array(img), np.array(txt_img)], axis=1)
return img
return img return image
def str_count(s): def str_count(s):
...@@ -213,7 +212,7 @@ def text_visual(texts, scores, img_h=400, img_w=600, threshold=0.): ...@@ -213,7 +212,7 @@ def text_visual(texts, scores, img_h=400, img_w=600, threshold=0.):
gap = font_size + 5 gap = font_size + 5
txt_img_list = [] txt_img_list = []
count, index = 0, 0 count, index = 1, 0
for idx, txt in enumerate(texts): for idx, txt in enumerate(texts):
index += 1 index += 1
if scores[idx] < threshold or math.isnan(scores[idx]): if scores[idx] < threshold or math.isnan(scores[idx]):
...@@ -230,24 +229,22 @@ def text_visual(texts, scores, img_h=400, img_w=600, threshold=0.): ...@@ -230,24 +229,22 @@ def text_visual(texts, scores, img_h=400, img_w=600, threshold=0.):
new_txt = ' ' + txt new_txt = ' ' + txt
draw_txt.text((0, gap * (count + 1)), new_txt, txt_color, font=font) draw_txt.text((0, gap * (count + 1)), new_txt, txt_color, font=font)
txt = tmp[img_w // font_size - 4:] txt = tmp[img_w // font_size - 4:]
count += 1
if count >= img_h // gap - 1: if count >= img_h // gap - 1:
txt_img_list.append(np.array(blank_img)) txt_img_list.append(np.array(blank_img))
blank_img, draw_txt = create_blank_img() blank_img, draw_txt = create_blank_img()
count = 0 count = 0
count += 1
if first_line: if first_line:
new_txt = str(index) + ': ' + txt + ' ' + '%.3f' % (scores[idx]) new_txt = str(index) + ': ' + txt + ' ' + '%.3f' % (scores[idx])
else: else:
new_txt = " " + txt + " " + '%.3f' % (scores[idx]) new_txt = " " + txt + " " + '%.3f' % (scores[idx])
draw_txt.text((0, gap * (count + 1)), new_txt, txt_color, font=font) draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
count += 1
# whether add new blank img or not # whether add new blank img or not
if count > img_h // gap - 1 and idx + 1 < len(texts): if count >= img_h // gap - 1 and idx + 1 < len(texts):
txt_img_list.append(np.array(blank_img)) txt_img_list.append(np.array(blank_img))
blank_img, draw_txt = create_blank_img() blank_img, draw_txt = create_blank_img()
count = 0 count = 0
count += 1
txt_img_list.append(np.array(blank_img)) txt_img_list.append(np.array(blank_img))
if len(txt_img_list) == 1: if len(txt_img_list) == 1:
blank_img = np.array(txt_img_list[0]) blank_img = np.array(txt_img_list[0])
...@@ -275,4 +272,4 @@ if __name__ == '__main__': ...@@ -275,4 +272,4 @@ if __name__ == '__main__':
new_img = draw_ocr(image, boxes, txts, scores, draw_txt=True) new_img = draw_ocr(image, boxes, txts, scores, draw_txt=True)
cv2.imwrite(img_name, new_img) cv2.imwrite(img_name, new_img)
\ No newline at end of file
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册