model_storage_directory参数取消，改为det_model_dir和rec_model_dir

paddleocr.py

@@ -29,25 +29,19 @@ from tools.infer import predict_system
 from ppocr.utils.utility import initial_logger
 
 logger = initial_logger()
-from ppocr.utils.utility import check_and_read_gif
+from ppocr.utils.utility import check_and_read_gif, get_image_file_list
 
 __all__ = ['PaddleOCR']
 
 model_params = {
-    'ch_det_mv3_db': {
-        'url':
-        'https://paddleocr.bj.bcebos.com/ch_models/ch_det_mv3_db_infer.tar',
-        'algorithm': 'DB',
-    },
-    'ch_rec_mv3_crnn_enhance': {
-        'url':
-        'https://paddleocr.bj.bcebos.com/ch_models/ch_rec_mv3_crnn_enhance_infer.tar',
-        'algorithm': 'CRNN'
-    },
+    'det': 'https://paddleocr.bj.bcebos.com/ch_models/ch_det_mv3_db_infer.tar',
+    'rec':
+    'https://paddleocr.bj.bcebos.com/ch_models/ch_rec_mv3_crnn_enhance_infer.tar',
 }
 
 SUPPORT_DET_MODEL = ['DB']
-SUPPORT_REC_MODEL = ['Rosetta', 'CRNN', 'STARNet', 'RARE']
+SUPPORT_REC_MODEL = ['CRNN']
+BASE_DIR = os.path.expanduser("~/.paddleocr/")
 
 
 def download_with_progressbar(url, save_path):
@@ -65,34 +59,29 @@ def download_with_progressbar(url, save_path):
         sys.exit(0)
 
 
-def download_and_unzip(url, model_storage_directory):
-    tmp_path = os.path.join(model_storage_directory, url.split('/')[-1])
-    print('download {} to {}'.format(url, tmp_path))
-    os.makedirs(model_storage_directory, exist_ok=True)
-    download_with_progressbar(url, tmp_path)
-    with tarfile.open(tmp_path, 'r') as tarObj:
-        for filename in tarObj.getnames():
-            tarObj.extract(filename, model_storage_directory)
-    os.remove(tmp_path)
-
-
-def maybe_download(model_storage_directory, model_name, mode='det'):
-    algorithm = None
+def maybe_download(model_storage_directory, url):
     # using custom model
-    if os.path.exists(os.path.join(model_name, 'model')) and os.path.exists(
-            os.path.join(model_name, 'params')):
-        return model_name, algorithm
-    # using the model of paddleocr
-    model_path = os.path.join(model_storage_directory, model_name)
-    if not os.path.exists(os.path.join(model_path,
-                                       'model')) or not os.path.exists(
-                                           os.path.join(model_path, 'params')):
-        assert model_name in model_params, 'model must in {}'.format(
-            model_params.keys())
-        download_and_unzip(model_params[model_name]['url'],
-                           model_storage_directory)
-        algorithm = model_params[model_name]['algorithm']
-    return model_path, algorithm
+    if not os.path.exists(os.path.join(
+            model_storage_directory, 'model')) or not os.path.exists(
+                os.path.join(model_storage_directory, 'params')):
+        tmp_path = os.path.join(model_storage_directory, url.split('/')[-1])
+        print('download {} to {}'.format(url, tmp_path))
+        os.makedirs(model_storage_directory, exist_ok=True)
+        download_with_progressbar(url, tmp_path)
+        with tarfile.open(tmp_path, 'r') as tarObj:
+            for member in tarObj.getmembers():
+                if "model" in member.name:
+                    filename = 'model'
+                elif "params" in member.name:
+                    filename = 'params'
+                else:
+                    continue
+                file = tarObj.extractfile(member)
+                with open(
+                        os.path.join(model_storage_directory, filename),
+                        'wb') as f:
+                    f.write(file.read())
+        os.remove(tmp_path)
 
 
 def parse_args():
@@ -111,7 +100,7 @@ def parse_args():
     # params for text detector
     parser.add_argument("--image_dir", type=str)
     parser.add_argument("--det_algorithm", type=str, default='DB')
-    parser.add_argument("--det_model_name", type=str, default='ch_det_mv3_db')
+    parser.add_argument("--det_model_dir", type=str, default=None)
     parser.add_argument("--det_max_side_len", type=float, default=960)
 
     # DB parmas
@@ -126,11 +115,11 @@ def parse_args():
 
     # params for text recognizer
     parser.add_argument("--rec_algorithm", type=str, default='CRNN')
-    parser.add_argument(
-        "--rec_model_name", type=str, default='ch_rec_mv3_crnn_enhance')
+    parser.add_argument("--rec_model_dir", type=str, default=None)
     parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320")
     parser.add_argument("--rec_char_type", type=str, default='ch')
     parser.add_argument("--rec_batch_num", type=int, default=30)
+    parser.add_argument("--max_text_length", type=int, default=25)
     parser.add_argument(
         "--rec_char_dict_path",
         type=str,
@@ -138,53 +127,30 @@ def parse_args():
     parser.add_argument("--use_space_char", type=bool, default=True)
     parser.add_argument("--enable_mkldnn", type=bool, default=False)
 
-    parser.add_argument("--model_storage_directory", type=str, default=False)
     parser.add_argument("--det", type=str2bool, default=True)
     parser.add_argument("--rec", type=str2bool, default=True)
     return parser.parse_args()
 
 
 class PaddleOCR(predict_system.TextSystem):
-    def __init__(self,
-                 det_model_name='ch_det_mv3_db',
-                 rec_model_name='ch_rec_mv3_crnn_enhance',
-                 model_storage_directory=None,
-                 log_level=20,
-                 **kwargs):
+    def __init__(self, **kwargs):
         """
         paddleocr package
         args:
-            det_model_name: det_model name, keep same with filename in paddleocr. default is ch_det_mv3_db
-            det_model_name: rec_model name, keep same with filename in paddleocr. default is ch_rec_mv3_crnn_enhance
-            model_storage_directory: model save path. default is ~/.paddleocr
-                                    det model will save to  model_storage_directory/det_model
-                                    rec model will save to  model_storage_directory/rec_model
-            log_level:
             **kwargs: other params show in paddleocr --help
         """
-        logger.setLevel(log_level)
         postprocess_params = parse_args()
-        # init model dir
-        if model_storage_directory:
-            self.model_storage_directory = model_storage_directory
-        else:
-            self.model_storage_directory = os.path.expanduser(
-                "~/.paddleocr/") + '/model'
-        Path(self.model_storage_directory).mkdir(parents=True, exist_ok=True)
+        postprocess_params.__dict__.update(**kwargs)
 
+        # init model dir
+        if postprocess_params.det_model_dir is None:
+            postprocess_params.det_model_dir = os.path.join(BASE_DIR, 'det')
+        if postprocess_params.rec_model_dir is None:
+            postprocess_params.rec_model_dir = os.path.join(BASE_DIR, 'rec')
+        print(postprocess_params)
         # download model
-        det_model_path, det_algorithm = maybe_download(
-            self.model_storage_directory, det_model_name, 'det')
-        rec_model_path, rec_algorithm = maybe_download(
-            self.model_storage_directory, rec_model_name, 'rec')
-        # update model and post_process params
-        postprocess_params.__dict__.update(**kwargs)
-        postprocess_params.det_model_dir = det_model_path
-        postprocess_params.rec_model_dir = rec_model_path
-        if det_algorithm is not None:
-            postprocess_params.det_algorithm = det_algorithm
-        if rec_algorithm is not None:
-            postprocess_params.rec_algorithm = rec_algorithm
+        maybe_download(postprocess_params.det_model_dir, model_params['det'])
+        maybe_download(postprocess_params.rec_model_dir, model_params['rec'])
 
         if postprocess_params.det_algorithm not in SUPPORT_DET_MODEL:
             logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
@@ -229,3 +195,18 @@ class PaddleOCR(predict_system.TextSystem):
                 img = [img]
             rec_res, elapse = self.text_recognizer(img)
             return rec_res
+
+
+def main():
+    # for com
+    args = parse_args()
+    image_file_list = get_image_file_list(args.image_dir)
+    if len(image_file_list) == 0:
+        logger.error('no images find in {}'.format(args.image_dir))
+        return
+    ocr_engine = PaddleOCR()
+    for img_path in image_file_list:
+        print(img_path)
+        result = ocr_engine.ocr(img_path, det=args.det, rec=args.rec)
+        for line in result:
+            print(line)