Optimize python inference part

python/README.md

@@ -1,36 +1,77 @@
-### 运行
+### 使用步骤
 1. 下载相应模型和用于显示的字体文件
-   - [提取码：30jv](https://pan.baidu.com/s/1qkqWK4wRdMjqGGbzR-FyWg)
-   - 下载之后放在`models`目录即可
-2. 运行`rapidOCR.py`文件, 即可
-   ```shell
-   cd python
-   python rapidOCR.py
-   ```
+   - [百度网盘 提取码：30jv](https://pan.baidu.com/s/1qkqWK4wRdMjqGGbzR-FyWg) | [Google Drive](https://drive.google.com/drive/folders/1x_a9KpCo_1blxH1xFOfgKVkw1HYRVywY?usp=sharing)
 
+   - 下载之后放在`models`目录即可，最终`models`目录结构如下：
+       ```text
+       models
+         |-- ch_PP-OCRv2_det_infer.onnx
+         |-- ch_ppocr_mobile_v2.0_cls_infer.onnx
+         |-- ch_ppocr_mobile_v2.0_det_infer.onnx
+         |-- ch_ppocr_server_v2.0_det_infer.onnx
+         |-- ch_ppocr_server_v2.0_rec_infer.onnx
+         |-- en_number_mobile_v2.0_rec_infer.onnx
+         |-- japan_rec_crnn.onnx
+         `-- msyh.ttc
+       ```
 
-### 文本检测+方向分类+文本识别
+2. 运行
+   - 接口调用方式运行
+     - 如果想要使用其他语言的识别模型，可以只对`rec_model_path`和`keys_path`作对应更改即可
+     - **!!!各个不同语言的推理代码区别只在于模型和字典文件!!!**
+        ```python
+        from rapid_ocr_api import TextSystem, visualize
 
-```python
-from ch_ppocr_mobile_v2_cls import TextClassifier
-from ch_ppocr_mobile_v2_det import TextDetector
-from ch_ppocr_mobile_v2_rec import TextRecognizer
+        det_model_path = 'models/ch_ppocr_mobile_v2.0_det_infer.onnx'
+        cls_model_path = 'models/ch_ppocr_mobile_v2.0_cls_infer.onnx'
 
-det_model_path = 'models/ch_ppocr_mobile_v2.0_det_infer.onnx'
-cls_model_path = 'models/ch_ppocr_mobile_v2.0_cls_infer.onnx'
-rec_model_path = 'models/ch_ppocr_mobile_v2.0_rec_infer.onnx'
+        # 中英文识别
+        rec_model_path = 'models/ch_ppocr_mobile_v2.0_rec_infer.onnx'
+        keys_path = 'ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt'
 
-image_path = r'test_images/det_images/1.jpg'
-keys_path = 'ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt'
+        text_sys = TextSystem(det_model_path,
+                            rec_model_path,
+                            use_angle_cls=True,
+                            cls_model_path=cls_model_path,
+                            keys_path=keys_path)
 
-text_sys = TextSystem(det_model_path,
-                     rec_model_path,
-                     use_angle_cls=True,
-                     cls_model_path=cls_model_path,
-                     keys_path=keys_path)
-dt_boxes, rec_res = text_sys(image_path)
-visualize(image_path, dt_boxes, rec_res)
-```
+        image_path = r'test_images/det_images/ch_en_num.jpg'
+        dt_boxes, rec_res = text_sys(image_path)
+        visualize(image_path, dt_boxes, rec_res)
+
+        # 只有中英文和数字识别
+        rec_model_path = 'models/en_number_mobile_v2.0_rec_infer.onnx'
+        keys_path = 'en_number_ppocr_mobile_v2_rec/en_dict.txt'
+
+        text_sys = TextSystem(det_model_path,
+                            rec_model_path,
+                            use_angle_cls=True,
+                            cls_model_path=cls_model_path,
+                            keys_path=keys_path)
+
+        image_path = r'test_images/det_images/en_num.png'
+        dt_boxes, rec_res = text_sys(image_path)
+        visualize(image_path, dt_boxes, rec_res)
+
+        # 日语识别
+        rec_model_path = 'models/japan_rec_crnn.onnx'
+        keys_path = 'japan_ppocr_mobile_v2_rec/japan_dict.txt'
+
+        text_sys = TextSystem(det_model_path,
+                            rec_model_path,
+                            use_angle_cls=True,
+                            cls_model_path=cls_model_path,
+                            keys_path=keys_path)
+
+        image_path = r'test_images/det_images/japan.png'
+        dt_boxes, rec_res = text_sys(image_path)
+        visualize(image_path, dt_boxes, rec_res)
+        ```
+
+    - 命令行运行
+        ```shell
+        $ bash test_demo.sh
+        ```
 
 ### 相关调节参数
 |参数名称|作用|建议取值范围|默认值|代码位置|

python/rapidOCR.sh

@@ -1,12 +0,0 @@
-#! /bin/bash
-
-# The img_path can be file path or folder path.
-img_path="test_images/det_images/1.jpg"
-
-keys_path="ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt"
-python rapidOCR.py --det_model_path models/ch_PP-OCRv2_det_infer.onnx \
-                   --cls_model_path models/ch_ppocr_mobile_v2.0_cls_infer.onnx \
-                   --rec_model_path models/ch_ppocr_mobile_v2.0_rec_infer.onnx \
-                   --image_path ${img_path} \
-                   --text_score 0.5 \
-                   --keys_path ${keys_path}

python/rapid_ocr_api.py

@@ -79,6 +79,10 @@ def check_and_read_gif(img_path):
 def draw_ocr_box_txt(image, boxes, txts,
                      scores=None, text_score=0.5,
                      font_path="./models/msyh.ttc"):
+    if not Path(font_path).exists():
+        raise FileNotFoundError(f'The {font_path} does not exists! \n'
+                                f'Please download the file in the https://drive.google.com/drive/folders/1x_a9KpCo_1blxH1xFOfgKVkw1HYRVywY')
+
     h, w = image.height, image.width
     img_left = image.copy()
     img_right = Image.new('RGB', (w, h), (255, 255, 255))
@@ -118,8 +122,7 @@ def draw_ocr_box_txt(image, boxes, txts,
                 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)
 
@@ -155,6 +158,8 @@ class TextSystem(object):
                  cls_model_path=None,
                  keys_path=None):
         super(TextSystem).__init__()
+        self.text_score = 0.5
+
         self.text_detector = TextDetector(det_model_path)
         self.text_recognizer = TextRecognizer(rec_model_path, keys_path)
         self.use_angle_cls = use_angle_cls
@@ -252,7 +257,7 @@ class TextSystem(object):
         filter_boxes, filter_rec_res = [], []
         for box, rec_reuslt in zip(dt_boxes, rec_res):
             text, score = rec_reuslt
-            if score >= text_score:
+            if score >= self.text_score:
                 filter_boxes.append(box)
                 filter_rec_res.append(rec_reuslt)
         return filter_boxes, filter_rec_res
@@ -270,12 +275,10 @@ if __name__ == '__main__':
 
     parser.add_argument('--image_path', type=str,
                         default='test_images/det_images/')
-    parser.add_argument('--text_score', type=float, default=0.5)
 
     parser.add_argument('--keys_path', type=str,
                         default="ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt")
     args = parser.parse_args()
-    text_score = args.text_score
 
     if args.det_model_path.find('server') != -1:
         from ch_ppocr_server_v2_det import TextDetector
@@ -287,6 +290,8 @@ if __name__ == '__main__':
     else:
         if args.rec_model_path.find('en') != -1:
             from en_number_ppocr_mobile_v2_rec import TextRecognizer
+        elif args.rec_model_path.find('japan') != -1:
+            from japan_ppocr_mobile_v2_rec import TextRecognizer
         else:
             from ch_ppocr_mobile_v2_rec import TextRecognizer
 

python/test_demo.py

@@ -0,0 +1,51 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from rapid_ocr_api import TextSystem, visualize
+
+det_model_path = 'models/ch_ppocr_mobile_v2.0_det_infer.onnx'
+cls_model_path = 'models/ch_ppocr_mobile_v2.0_cls_infer.onnx'
+
+# 中英文识别
+rec_model_path = 'models/ch_ppocr_mobile_v2.0_rec_infer.onnx'
+keys_path = 'ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt'
+
+text_sys = TextSystem(det_model_path,
+                      rec_model_path,
+                      use_angle_cls=True,
+                      cls_model_path=cls_model_path,
+                      keys_path=keys_path)
+
+image_path = r'test_images/det_images/ch_en_num.jpg'
+dt_boxes, rec_res = text_sys(image_path)
+visualize(image_path, dt_boxes, rec_res)
+
+
+# 只有中英文和数字识别
+rec_model_path = 'models/en_number_mobile_v2.0_rec_infer.onnx'
+keys_path = 'en_number_ppocr_mobile_v2_rec/en_dict.txt'
+
+text_sys = TextSystem(det_model_path,
+                      rec_model_path,
+                      use_angle_cls=True,
+                      cls_model_path=cls_model_path,
+                      keys_path=keys_path)
+
+image_path = r'test_images/det_images/en_num.png'
+dt_boxes, rec_res = text_sys(image_path)
+visualize(image_path, dt_boxes, rec_res)
+
+
+# 日语识别
+rec_model_path = 'models/japan_rec_crnn.onnx'
+keys_path = 'japan_ppocr_mobile_v2_rec/japan_dict.txt'
+
+text_sys = TextSystem(det_model_path,
+                      rec_model_path,
+                      use_angle_cls=True,
+                      cls_model_path=cls_model_path,
+                      keys_path=keys_path)
+
+image_path = r'test_images/det_images/japan.png'
+dt_boxes, rec_res = text_sys(image_path)
+visualize(image_path, dt_boxes, rec_res)

python/test_demo.sh

@@ -0,0 +1,16 @@
+#! /bin/bash
+
+det_model_path="models/ch_PP-OCRv2_det_infer.onnx"
+cls_model_path="models/ch_ppocr_mobile_v2.0_cls_infer.onnx"
+
+# 这里可以按照自己需求，做相关更改
+rec_model_path="models/ch_ppocr_mobile_v2.0_rec_infer.onnx"
+keys_path="ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt"
+
+img_path="test_images/det_images/ch_en_num.jpg"
+
+python rapid_ocr_api.py --det_model_path ${det_model_path} \
+                        --cls_model_path ${cls_model_path} \
+                        --rec_model_path ${rec_model_path} \
+                        --image_path ${img_path} \
+                        --keys_path ${keys_path}