Merge pull request #99 from LDOUBLEV/fixocr

opt visualized func and add docker usage in cpu

doc/installation.md

@@ -15,6 +15,9 @@ cd /home/Projects
 # 首次运行需创建一个docker容器，再次运行时不需要运行当前命令
 # 创建一个名字为ppocr的docker容器，并将当前目录映射到容器的/paddle目录下
 
+如果您希望在CPU环境下使用docker，使用docker而不是nvidia-docker创建docker
+sudo docker run --name ppocr -v $PWD:/paddle --network=host -it hub.baidubce.com/paddlepaddle/paddle:latest-gpu-cuda9.0-cudnn7-dev /bin/bash
+
 如果您的机器安装的是CUDA9，请运行以下命令创建容器
 sudo nvidia-docker run --name ppocr -v $PWD:/paddle --network=host -it hub.baidubce.com/paddlepaddle/paddle:latest-gpu-cuda9.0-cudnn7-dev /bin/bash
 
@@ -24,7 +27,7 @@ sudo nvidia-docker run --name ppocr -v $PWD:/paddle --network=host -it hub.baidu
 您也可以访问[DockerHub](https://hub.docker.com/r/paddlepaddle/paddle/tags/)获取与您机器适配的镜像。
 
 # ctrl+P+Q可退出docker，重新进入docker使用如下命令
-sudo nvidia-docker container exec -it ppocr /bin/bash
+sudo docker container exec -it ppocr /bin/bash
 ```
 
 注意：如果docker pull过慢，可以按照如下步骤手动下载后加载docker,以cuda9 docker为例，使用cuda10 docker只需要将cuda9改为cuda10即可。

tools/infer/utility.py

@@ -23,6 +23,7 @@ import cv2
 import numpy as np
 import json
 from PIL import Image, ImageDraw, ImageFont
+import math
 
 
 def parse_args():
@@ -127,6 +128,18 @@ def resize_img(img, input_size=600):
 
 
 def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
+    """
+    Visualize the results of OCR detection and recognition
+    args:
+        image(Image): image from Image.open
+        boxes(list): boxes with shape(N, 4, 2)
+        txts(list): the texts
+        scores(list): txxs corresponding scores
+        draw_txt(bool): whether draw text or not
+        drop_score(float): only scores greater than drop_threshold will be visualized
+    return(array):
+        the visualized img
+    """
     from PIL import Image, ImageDraw, ImageFont
 
     img = image.copy()
@@ -154,32 +167,111 @@ def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
             fill='red')
 
     if draw_txt:
-        txt_color = (0, 0, 0)
-        img = np.array(resize_img(img))
-        _h = img.shape[0]
-        blank_img = np.ones(shape=[_h, 600], dtype=np.int8) * 255
+        img = np.array(resize_img(img, input_size=600))
+        txt_img = text_visual(
+            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)
+
+    return img
+
+
+def str_count(s):
+    """
+    Count the number of Chinese characters,
+    a single English character and a single number
+    equal to half the length of Chinese characters.
+
+    args:
+        s(string): the input of string
+    return(int):
+        the number of Chinese characters
+    """
+    import string
+    count_zh = count_pu = 0
+    s_len = len(s)
+    en_dg_count = 0
+    for c in s:
+        if c in string.ascii_letters or c.isdigit() or c.isspace():
+            en_dg_count += 1
+        elif c.isalpha():
+            count_zh += 1
+        else:
+            count_pu += 1
+    return s_len - math.ceil(en_dg_count / 2)
+
+
+def text_visual(texts, scores, img_h=400, img_w=600, threshold=0.):
+    """
+    create new blank img and draw txt on it
+    args:
+        texts(list): the text will be draw
+        scores(list|None): corresponding score of each txt
+        img_h(int): the height of blank img
+        img_w(int): the width of blank img
+    return(array):
+
+    """
+    if scores is not None:
+        assert len(texts) == len(
+            scores), "The number of txts and corresponding scores must match"
+
+    def create_blank_img():
+        blank_img = np.ones(shape=[img_h, img_w], dtype=np.int8) * 255
+        blank_img[:, img_w - 1:] = 0
         blank_img = Image.fromarray(blank_img).convert("RGB")
         draw_txt = ImageDraw.Draw(blank_img)
+        return blank_img, draw_txt
+
+    blank_img, draw_txt = create_blank_img()
+
+    font_size = 20
+    txt_color = (0, 0, 0)
+    font = ImageFont.truetype("./doc/simfang.ttf", font_size, encoding="utf-8")
 
-        font_size = 20
-        gap = 20
-        title = "index           text           score"
-        font = ImageFont.truetype(
-            "./doc/simfang.ttf", font_size, encoding="utf-8")
-
-        draw_txt.text((20, 0), title, txt_color, font=font)
-        count = 0
-        for idx, txt in enumerate(txts):
-            if scores[idx] < drop_score:
-                continue
-            font = ImageFont.truetype(
-                "./doc/simfang.ttf", font_size, encoding="utf-8")
-            new_txt = str(idx) + ':  ' + txt + '    ' + '%.3f' % (scores[idx])
-            draw_txt.text(
-                (20, gap * (count + 1)), new_txt, txt_color, font=font)
+    gap = font_size + 5
+    txt_img_list = []
+    count, index = 0, 0
+    for idx, txt in enumerate(texts):
+        index += 1
+        if scores[idx] < threshold:
+            index -= 1
+            continue
+        first_line = True
+        while str_count(txt) >= img_w // font_size - 4:
+            tmp = txt
+            txt = tmp[:img_w // font_size - 4]
+            if first_line:
+                new_txt = str(index) + ': ' + txt
+                first_line = False
+            else:
+                new_txt = '    ' + txt
+            draw_txt.text((0, gap * (count + 1)), new_txt, txt_color, font=font)
+            txt = tmp[img_w // font_size - 4:]
             count += 1
-        img = np.concatenate([np.array(img), np.array(blank_img)], axis=1)
-    return img
+            if count >= img_h // gap - 1:
+                txt_img_list.append(np.array(blank_img))
+                blank_img, draw_txt = create_blank_img()
+                count = 0
+
+        if first_line:
+            new_txt = str(index) + ': ' + txt + '   ' + '%.3f' % (scores[idx])
+        else:
+            new_txt = "   " + txt + "  " + '%.3f' % (scores[idx])
+        draw_txt.text((0, gap * (count + 1)), new_txt, txt_color, font=font)
+        count += 1
+        # whether add new blank img or not
+        if count >= img_h // gap - 1 and idx + 1 < len(texts):
+            txt_img_list.append(np.array(blank_img))
+            blank_img, draw_txt = create_blank_img()
+            count = 0
+
+    txt_img_list.append(np.array(blank_img))
+    if len(txt_img_list) == 1:
+        blank_img = np.array(txt_img_list[0])
+    else:
+        blank_img = np.concatenate(txt_img_list, axis=1)
+    # cv2.imwrite("./draw_txt.jpg", np.array(blank_img))
+    return np.array(blank_img)
 
 
 if __name__ == '__main__':