Update models

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.TTF filter=lfs diff=lfs merge=lfs -text
 images/car_plate.jpeg filter=lfs diff=lfs merge=lfs -text
+*.ttc filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,87 +1,19 @@
 # -*- encoding: utf-8 -*-
-import math
-import random
 import time
 from pathlib import Path
 
 import cv2
 import gradio as gr
-import numpy as np
-from PIL import Image, ImageDraw, ImageFont
 from rapidocr_onnxruntime import RapidOCR
 
+from utils import visualize
 
-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="./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
+font_dict = {
+    'ch': 'FZYTK.TTF',
+    'japan': 'japan.ttc',
+    'korean': 'korean.ttf',
+    'en': 'FZYTK.TTF'
+}
 
 
 def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5,
@@ -102,8 +34,17 @@ def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5,
                          rec_img_shape=rec_image_shape)
     elapse = time.time() - s
 
+    if 'ch' in rec_model_path or 'en' in rec_model_path:
+        lan_name = 'ch'
+    elif 'japan' in rec_model_path:
+        lan_name = 'japan'
+    elif 'korean' in rec_model_path:
+        lan_name = 'korean'
+    else:
+        lan_name = 'ch'
+
     out_log_list.append(f'Init Model cost: {elapse:.5f}')
-    out_log_list.extend([f'det_model:{det_model_path}',
+    out_log_list.extend([f'det_model: {det_model_path}',
                          f'rec_model: {rec_model_path}',
                          f'rec_image_shape: {rec_image_shape}'])
 
@@ -120,73 +61,78 @@ def inference(img_path, box_thresh=0.5, unclip_ratio=1.6, text_score=0.5,
         return img_path, '未识别到有效文本', out_log
 
     dt_boxes, rec_res, scores = list(zip(*ocr_result))
-    img_save_path = visualize(img_path, dt_boxes, rec_res, scores)
+    font_path = Path('fonts') / font_dict.get(lan_name)
+    img_save_path = visualize(img_path, dt_boxes, rec_res, scores,
+                              font_path=str(font_path))
     output_text = [f'{one_rec} {float(score):.4f}'
                    for one_rec, score in zip(rec_res, scores)]
     return img_save_path, output_text, out_log
 
 
-examples = [['images/1.jpg'], 
-            ['images/ch_en_num.jpg'], 
-            ['images/air_ticket.jpg'], 
-            ['images/car_plate.jpeg'], 
-            ['images/idcard.jpg'], 
-            ['images/train_ticket.jpeg']]
-
-with gr.Blocks(title='RapidOCR') as demo:
-    gr.Markdown("""
-        <h1><center><a href="https://github.com/RapidAI/RapidOCR" target="_blank">Rapid⚡OCR</a></center></h1>
-
-        ### Docs: [Docs](https://rapidocr.rtfd.io/)
-        ### Parameters docs: [link](https://github.com/RapidAI/RapidOCR/tree/main/python#configyaml%E4%B8%AD%E5%B8%B8%E7%94%A8%E5%8F%82%E6%95%B0%E4%BB%8B%E7%BB%8D)
-        - **box_thresh**: 检测到的框是文本的概率，值越大，框中是文本的概率就越大。存在漏检时，调低该值。取值范围：[0, 1.0]
-        - **unclip_ratio**: 控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]
-        - **text_score**: 文本识别结果是正确的置信度，值越大，显示出的识别结果更准确。存在漏检时，调低该值。取值范围：[0, 1.0]
-        ### 运行环境：
-        Python: 3.8 | onnxruntime: 1.14.1 | rapidocr_onnxruntime: 1.2.5""")
-    gr.Markdown('**超参数调节**')
-    with gr.Row():
-        box_thresh = gr.Slider(minimum=0, maximum=1.0, value=0.5,
-                               label='box_thresh', step=0.1,
-                               interactive=True,
-                               info='[0, 1.0]')
-        unclip_ratio = gr.Slider(minimum=1.5, maximum=2.0, value=1.6,
-                                 label='unclip_ratio', step=0.1,
-                                 interactive=True,
-                                 info='[1.5, 2.0]')
-        text_score = gr.Slider(minimum=0, maximum=1.0, value=0.5,
-                               label='text_score', step=0.1,
-                               interactive=True,
-                               info='[0, 1.0]')
-
-    gr.Markdown('**模型选择**')
-    with gr.Row():
-        text_det = gr.Dropdown(['ch_PP-OCRv3_det_infer.onnx',
-                                'ch_PP-OCRv2_det_infer.onnx',
-                                'ch_ppocr_server_v2.0_det_infer.onnx'],
-                               label='选择文本检测模型',
-                               value='ch_PP-OCRv3_det_infer.onnx',
-                               interactive=True)
-
-        text_rec = gr.Dropdown(['ch_PP-OCRv3_rec_infer.onnx',
-                                'ch_PP-OCRv2_rec_infer.onnx',
-                                'ch_ppocr_server_v2.0_rec_infer.onnx'],
-                               label='选择文本识别模型',
-                               value='ch_PP-OCRv3_rec_infer.onnx',
-                               interactive=True)
-
-    with gr.Row():
-        input_img = gr.Image(type='filepath', label='Input')
-        out_img = gr.Image(type='filepath', label='Output')
-    out_log = gr.outputs.Textbox(type='text', label='Run Log')
-    out_txt = gr.outputs.Textbox(type='text', label='RecText')
-    button = gr.Button('Submit')
-    button.click(fn=inference,
-                 inputs=[input_img, box_thresh, unclip_ratio, text_score,
-                         text_det, text_rec],
-                 outputs=[out_img, out_txt, out_log])
-    gr.Examples(examples=examples,
-                inputs=[input_img, box_thresh, unclip_ratio, text_score,
-                        text_det, text_rec],
-                outputs=[out_img, out_txt, out_log], fn=inference)
-demo.launch(debug=True, enable_queue=True)
+if __name__ == '__main__':
+
+    examples = [['images/1.jpg'],
+                ['images/ch_en_num.jpg'],
+                ['images/air_ticket.jpg'],
+                ['images/car_plate.jpeg'],
+                ['images/idcard.jpg'],
+                ['images/train_ticket.jpeg'],
+                ['images/japan_2.jpg'],
+                ['images/korean_1.jpg']]
+
+    with gr.Blocks(title='RapidOCR') as demo:
+        gr.Markdown("""
+            <h1><center><a href="https://github.com/RapidAI/RapidOCR" target="_blank">Rapid⚡OCR</a></center></h1>
+
+            ### Docs: [Docs](https://rapidocr.rtfd.io/)
+            ### 运行环境：
+            Python: 3.8 | onnxruntime: 1.14.1 | rapidocr_onnxruntime: 1.2.5""")
+        gr.Markdown(
+            '''**[超参数调节](https://github.com/RapidAI/RapidOCR/tree/main/python#configyaml%E4%B8%AD%E5%B8%B8%E7%94%A8%E5%8F%82%E6%95%B0%E4%BB%8B%E7%BB%8D)**
+            - **box_thresh**: 检测到的框是文本的概率，值越大，框中是文本的概率就越大。存在漏检时，调低该值。取值范围：[0, 1.0]
+            - **unclip_ratio**: 控制文本检测框的大小，值越大，检测框整体越大。在出现框截断文字的情况，调大该值。取值范围：[1.5, 2.0]
+            - **text_score**: 文本识别结果是正确的置信度，值越大，显示出的识别结果更准确。存在漏检时，调低该值。取值范围：[0, 1.0]
+            ''')
+        with gr.Row():
+            box_thresh = gr.Slider(minimum=0, maximum=1.0, value=0.5,
+                                   label='box_thresh', step=0.1,
+                                   interactive=True,
+                                   info='[0, 1.0]')
+            unclip_ratio = gr.Slider(minimum=1.5, maximum=2.0, value=1.6,
+                                     label='unclip_ratio', step=0.1,
+                                     interactive=True,
+                                     info='[1.5, 2.0]')
+            text_score = gr.Slider(minimum=0, maximum=1.0, value=0.5,
+                                   label='text_score', step=0.1,
+                                   interactive=True,
+                                   info='[0, 1.0]')
+
+        gr.Markdown('**[模型选择](https://github.com/RapidAI/RapidOCR/blob/main/docs/models.md)**')
+        with gr.Row():
+            text_det = gr.Dropdown(['ch_PP-OCRv3_det_infer.onnx',
+                                    'ch_PP-OCRv2_det_infer.onnx',
+                                    'ch_ppocr_server_v2.0_det_infer.onnx'],
+                                   label='选择文本检测模型',
+                                   value='ch_PP-OCRv3_det_infer.onnx',
+                                   interactive=True)
+            rec_model_list = [v.name for v in Path('models/text_rec').iterdir()]
+            text_rec = gr.Dropdown(rec_model_list,
+                                   label='选择文本识别模型(包括中英文和多语言)',
+                                   value='ch_PP-OCRv3_rec_infer.onnx',
+                                   interactive=True)
+
+        with gr.Row():
+            input_img = gr.Image(type='filepath', label='Input')
+            out_img = gr.Image(type='filepath', label='Output')
+        out_log = gr.outputs.Textbox(type='text', label='Run Log')
+        out_txt = gr.outputs.Textbox(type='text', label='RecText')
+        button = gr.Button('Submit')
+        button.click(fn=inference,
+                     inputs=[input_img, box_thresh, unclip_ratio, text_score,
+                             text_det, text_rec],
+                     outputs=[out_img, out_txt, out_log])
+        gr.Examples(examples=examples,
+                    inputs=[input_img, box_thresh, unclip_ratio, text_score,
+                            text_det, text_rec],
+                    outputs=[out_img, out_txt, out_log], fn=inference)
+    demo.launch(debug=True, enable_queue=True)

FZYTK.TTF → models/text_rec/en_PP-OCRv3_rec_infer.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4065a23df6823c8e2b69a0e76d02f02a6470b8774a5e91086609701ad95cc33f
-size 3241748
+oid sha256:ef7abd8bd3629ae57ea2c28b425c1bd258a871b93fd2fe7c433946ade9b5d9ea
+size 8967018

models/text_rec/en_number_mobile_v2.0_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e679ba625c544444be78292a50d9e1af9caa1569239a88bb8b864cb688b11c01
+size 1882607

models/text_rec/japan_rec_crnn_v2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b0495059f5738166e606d864b04ff00093f67a807efb02cddf472839cae970c
+size 3571807

models/text_rec/korean_mobile_v2.0_rec_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6558500138b43b46a4941957fb8c918546dae5fb0e71718536f1883acc80faf
+size 3290650