Fix bugs of inputting paramters with class instance

python/demo.py

@@ -93,7 +93,6 @@ if __name__ == '__main__':
     image_path = 'tests/test_files/ch_en_num.jpg'
     with open(image_path, 'rb') as f:
         img = f.read()
-        # img = cv2.imread(image_path)
     result, elapse_list = rapid_ocr(img)
     print(result)
     print(elapse_list)

python/rapidocr_onnxruntime/rapid_ocr_api.py

@@ -10,7 +10,8 @@ from typing import Union
 import cv2
 import numpy as np
 
-from .utils import LoadImage, ParseArgs, concat_model_path, read_yaml
+from .utils import (LoadImage, init_args, UpdateParameters, concat_model_path,
+                    read_yaml)
 
 root_dir = Path(__file__).resolve().parent
 sys.path.append(str(root_dir))
@@ -25,8 +26,8 @@ class RapidOCR():
         config = concat_model_path(config)
 
         if kwargs:
-            parser = ParseArgs()
-            config = parser.update_config(config, **kwargs)
+            updater = UpdateParameters()
+            config = updater(config, **kwargs)
 
         global_config = config['Global']
         self.print_verbose = global_config['print_verbose']
@@ -180,7 +181,7 @@ class RapidOCR():
 
 
 def main():
-    args = ParseArgs().args
+    args = init_args()
     ocr_engine = RapidOCR(**vars(args))
 
     result, elapse_list = ocr_engine(args.img_path)

python/rapidocr_onnxruntime/utils.py

@@ -133,7 +133,7 @@ class LoadImage():
         '''RGBA → RGB
         '''
         r, g, b, a = cv2.split(img)
-        new_img  = cv2.merge((b, g, r))
+        new_img = cv2.merge((b, g, r))
 
         not_a = cv2.bitwise_not(a)
         not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR)
@@ -166,55 +166,55 @@ def concat_model_path(config):
     return config
 
 
-class ParseArgs():
-    def __init__(self, ):
-        self.args = self.init_args()
-        self.args_dict = vars(self.args)
+def init_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-img', '--img_path', type=str, default=None,
+                        required=True)
+    parser.add_argument('-p', '--print_cost',
+                        action='store_true', default=False)
 
-    def init_args(self):
-        parser = argparse.ArgumentParser()
-        parser.add_argument('-img', '--img_path', type=str, default=None,
-                            required=True)
-        parser.add_argument('-p', '--print_cost',
-                            action='store_true', default=False)
+    global_group = parser.add_argument_group(title='Global')
+    global_group.add_argument('--text_score', type=float, default=0.5)
+    global_group.add_argument('--use_angle_cls', type=bool, default=True)
+    global_group.add_argument('--use_text_det', type=bool, default=True)
+    global_group.add_argument('--print_verbose', type=bool, default=False)
+    global_group.add_argument('--min_height', type=int, default=30)
+    global_group.add_argument('--width_height_ratio', type=int, default=8)
 
-        global_group = parser.add_argument_group(title='Global')
-        global_group.add_argument('--text_score', type=float, default=0.5)
-        global_group.add_argument('--use_angle_cls', type=bool, default=True)
-        global_group.add_argument('--use_text_det', type=bool, default=True)
-        global_group.add_argument('--print_verbose', type=bool, default=False)
-        global_group.add_argument('--min_height', type=int, default=30)
-        global_group.add_argument('--width_height_ratio', type=int, default=8)
+    det_group = parser.add_argument_group(title='Det')
+    det_group.add_argument('--det_model_path', type=str, default=None)
+    det_group.add_argument('--det_limit_side_len', type=float, default=736)
+    det_group.add_argument('--det_limit_type', type=str, default='min',
+                           choices=['max', 'min'])
+    det_group.add_argument('--det_thresh', type=float, default=0.3)
+    det_group.add_argument('--det_box_thresh', type=float, default=0.5)
+    det_group.add_argument('--det_unclip_ratio', type=float, default=1.6)
+    det_group.add_argument('--det_use_dilation', type=bool, default=True)
+    det_group.add_argument('--det_score_mode', type=str, default='fast',
+                           choices=['slow', 'fast'])
 
-        det_group = parser.add_argument_group(title='Det')
-        det_group.add_argument('--det_model_path', type=str, default=None)
-        det_group.add_argument('--det_limit_side_len', type=float, default=736)
-        det_group.add_argument('--det_limit_type', type=str, default='min',
-                               choices=['max', 'min'])
-        det_group.add_argument('--det_thresh', type=float, default=0.3)
-        det_group.add_argument('--det_box_thresh', type=float, default=0.5)
-        det_group.add_argument('--det_unclip_ratio', type=float, default=1.6)
-        det_group.add_argument('--det_use_dilation', type=bool, default=True)
-        det_group.add_argument('--det_score_mode', type=str, default='fast',
-                               choices=['slow', 'fast'])
+    cls_group = parser.add_argument_group(title='Cls')
+    cls_group.add_argument('--cls_model_path', type=str, default=None)
+    cls_group.add_argument('--cls_image_shape', type=list,
+                           default=[3, 48, 192])
+    cls_group.add_argument('--cls_label_list', type=list,
+                           default=['0', '180'])
+    cls_group.add_argument('--cls_batch_num', type=int, default=6)
+    cls_group.add_argument('--cls_thresh', type=float, default=0.9)
 
-        cls_group = parser.add_argument_group(title='Cls')
-        cls_group.add_argument('--cls_model_path', type=str, default=None)
-        cls_group.add_argument('--cls_image_shape', type=list,
-                               default=[3, 48, 192])
-        cls_group.add_argument('--cls_label_list', type=list,
-                               default=['0', '180'])
-        cls_group.add_argument('--cls_batch_num', type=int, default=6)
-        cls_group.add_argument('--cls_thresh', type=float, default=0.9)
+    rec_group = parser.add_argument_group(title='Rec')
+    rec_group.add_argument('--rec_model_path', type=str, default=None)
+    rec_group.add_argument('--rec_image_shape', type=list,
+                           default=[3, 48, 320])
+    rec_group.add_argument('--rec_batch_num', type=int, default=6)
 
-        rec_group = parser.add_argument_group(title='Rec')
-        rec_group.add_argument('--rec_model_path', type=str, default=None)
-        rec_group.add_argument('--rec_image_shape', type=list,
-                               default=[3, 48, 320])
-        rec_group.add_argument('--rec_batch_num', type=int, default=6)
+    args = parser.parse_args()
+    return args
 
-        args = parser.parse_args()
-        return args
+
+class UpdateParameters():
+    def __init__(self) -> None:
+        pass
 
     def parse_kwargs(self, **kwargs):
         global_dict, det_dict, cls_dict, rec_dict = {}, {}, {}, {}
@@ -229,7 +229,7 @@ class ParseArgs():
                 global_dict[k] = v
         return global_dict, det_dict, cls_dict, rec_dict
 
-    def update_config(self, config, **kwargs):
+    def __call__(self, config, **kwargs):
         global_dict, det_dict, cls_dict, rec_dict = self.parse_kwargs(**kwargs)
         new_config = {
             'Global': self.update_global_params(config['Global'],
@@ -241,38 +241,44 @@ class ParseArgs():
         return new_config
 
     def update_global_params(self, config, global_dict):
-        config.update(global_dict)
+        if global_dict:
+            config.update(global_dict)
         return config
 
     def update_det_params(self, config, det_dict):
-        det_dict = {k.split('det_')[1]: v for k, v in det_dict.items()}
-        if not det_dict['model_path']:
-            det_dict['model_path'] = str(root_dir / config['model_path'])
-        config.update(det_dict)
+        if det_dict:
+            det_dict = {k.split('det_')[1]: v for k, v in det_dict.items()}
+            if not det_dict['model_path']:
+                det_dict['model_path'] = str(root_dir / config['model_path'])
+            config.update(det_dict)
         return config
 
     def update_cls_params(self, config, cls_dict):
-        need_remove_prefix = ['cls_label_list', 'cls_model_path']
-        new_cls_dict = {}
-        for k, v in cls_dict.items():
-            if k in need_remove_prefix:
-                k = k.split('cls_')[1]
-            new_cls_dict[k] = v
+        if cls_dict:
+            need_remove_prefix = ['cls_label_list', 'cls_model_path']
+            new_cls_dict = {}
+            for k, v in cls_dict.items():
+                if k in need_remove_prefix:
+                    k = k.split('cls_')[1]
+                new_cls_dict[k] = v
 
-        if not new_cls_dict['model_path']:
-            new_cls_dict['model_path'] = str(root_dir / config['model_path'])
-        config.update(new_cls_dict)
+            if not new_cls_dict['model_path']:
+                new_cls_dict['model_path'] = str(
+                    root_dir / config['model_path'])
+            config.update(new_cls_dict)
         return config
 
     def update_rec_params(self, config, rec_dict):
-        need_remove_prefix = ['rec_model_path']
-        new_rec_dict = {}
-        for k, v in rec_dict.items():
-            if k in need_remove_prefix:
-                k = k.split('rec_')[1]
-            new_rec_dict[k] = v
+        if rec_dict:
+            need_remove_prefix = ['rec_model_path']
+            new_rec_dict = {}
+            for k, v in rec_dict.items():
+                if k in need_remove_prefix:
+                    k = k.split('rec_')[1]
+                new_rec_dict[k] = v
 
-        if not new_rec_dict['model_path']:
-            new_rec_dict['model_path'] = str(root_dir / config['model_path'])
-        config.update(new_rec_dict)
+            if not new_rec_dict['model_path']:
+                new_rec_dict['model_path'] = str(
+                    root_dir / config['model_path'])
+            config.update(new_rec_dict)
         return config

python/tests/test_all_ort.py

@@ -61,3 +61,38 @@ def test_input_path():
     result, _ = rapid_ocr(image_path)
     assert result[0][1] == '正品促销'
     assert len(result) == 17
+
+
+def test_input_parameters():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    rapid_ocr = RapidOCR(text_score=1)
+    result, _ = rapid_ocr(image_path)
+
+    assert result is None
+
+
+def test_input_det_parameters():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    with pytest.raises(FileNotFoundError) as exc_info:
+        rapid_ocr = RapidOCR(det_model_path='1.onnx')
+        result, _ = rapid_ocr(image_path)
+        raise FileNotFoundError()
+    assert exc_info.type is FileNotFoundError
+
+
+def test_input_cls_parameters():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    with pytest.raises(FileNotFoundError) as exc_info:
+        rapid_ocr = RapidOCR(cls_model_path='1.onnx')
+        result, _ = rapid_ocr(image_path)
+        raise FileNotFoundError()
+    assert exc_info.type is FileNotFoundError
+
+
+def test_input_rec_parameters():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    with pytest.raises(FileNotFoundError) as exc_info:
+        rapid_ocr = RapidOCR(rec_model_path='1.onnx')
+        result, _ = rapid_ocr(image_path)
+        raise FileNotFoundError()
+    assert exc_info.type is FileNotFoundError