Update files

utils.py

@@ -0,0 +1,85 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: [email protected]
+import math
+import random
+import time
+from pathlib import Path
+
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+
+def draw_ocr_box_txt(image, boxes, txts, font_path,
+                     scores=None, text_score=0.5):
+    h, w = image.height, image.width
+    img_left = image.copy()
+    img_right = Image.new('RGB', (w, h), (255, 255, 255))
+
+    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 float(scores[idx]) < text_score:
+            continue
+
+        color = (random.randint(0, 255),
+                 random.randint(0, 255),
+                 random.randint(0, 255))
+
+        box = [tuple(v) for v in box]
+        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 visualize(image_path, boxes, txts, scores,
+              font_path="./fonts/FZYTK.TTF"):
+    image = Image.open(image_path)
+
+    draw_img = draw_ocr_box_txt(image, boxes,
+                                txts, font_path,
+                                scores,
+                                text_score=0.5)
+
+    draw_img_save = Path("./inference_results/")
+    if not draw_img_save.exists():
+        draw_img_save.mkdir(parents=True, exist_ok=True)
+
+    time_stamp = time.strftime(
+        '%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
+    image_save = str(draw_img_save / f'{time_stamp}_{Path(image_path).name}')
+    cv2.imwrite(image_save, draw_img[:, :, ::-1])
+    return image_save