paddleocr.py 24.6 KB
Newer Older
WenmuZhou's avatar
WenmuZhou committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
WenmuZhou's avatar
WenmuZhou committed
17
import importlib
WenmuZhou's avatar
WenmuZhou committed
18
19

__dir__ = os.path.dirname(__file__)
zhoujun's avatar
zhoujun committed
20
21
22

import paddle

WenmuZhou's avatar
WenmuZhou committed
23
24
25
sys.path.append(os.path.join(__dir__, ''))

import cv2
WenmuZhou's avatar
WenmuZhou committed
26
import logging
WenmuZhou's avatar
WenmuZhou committed
27
28
29
import numpy as np
from pathlib import Path

WenmuZhou's avatar
WenmuZhou committed
30
31
32
33
tools = importlib.import_module('.', 'tools')
ppocr = importlib.import_module('.', 'ppocr')
ppstructure = importlib.import_module('.', 'ppstructure')

WenmuZhou's avatar
WenmuZhou committed
34
from tools.infer import predict_system
WenmuZhou's avatar
WenmuZhou committed
35
from ppocr.utils.logging import get_logger
WenmuZhou's avatar
WenmuZhou committed
36

WenmuZhou's avatar
WenmuZhou committed
37
logger = get_logger()
38
from ppocr.utils.utility import check_and_read_gif, get_image_file_list
39
from ppocr.utils.network import maybe_download, download_with_progressbar, is_link, confirm_model_dir_url
zhoujun's avatar
zhoujun committed
40
from tools.infer.utility import draw_ocr, str2bool, check_gpu
41
from ppstructure.utility import init_args, draw_structure_result
42
from ppstructure.predict_system import StructureSystem, save_structure_res
WenmuZhou's avatar
WenmuZhou committed
43

WenmuZhou's avatar
WenmuZhou committed
44
45
46
47
48
49
__all__ = [
    'PaddleOCR', 'PPStructure', 'draw_ocr', 'draw_structure_result',
    'save_structure_res', 'download_with_progressbar'
]

SUPPORT_DET_MODEL = ['DB']
50
VERSION = '2.5.0.1'
WenmuZhou's avatar
WenmuZhou committed
51
52
53
SUPPORT_REC_MODEL = ['CRNN']
BASE_DIR = os.path.expanduser("~/.paddleocr/")

54
55
56
57
DEFAULT_OCR_MODEL_VERSION = 'PP-OCRv3'
SUPPORT_OCR_MODEL_VERSION = ['PP-OCR', 'PP-OCRv2', 'PP-OCRv3']
DEFAULT_STRUCTURE_MODEL_VERSION = 'PP-STRUCTURE'
SUPPORT_STRUCTURE_MODEL_VERSION = ['PP-STRUCTURE']
WenmuZhou's avatar
WenmuZhou committed
58
MODEL_URLS = {
zhoujun's avatar
zhoujun committed
59
    'OCR': {
60
61
62
63
64
65
66
67
68
69
        'PP-OCRv3': {
            'det': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_det_infer.tar',
                },
                'en': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar',
                },
andyjpaddle's avatar
andyjpaddle committed
70
71
72
73
                'ml': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/Multilingual_PP-OCRv3_det_infer.tar'
                }
74
75
76
77
78
79
80
81
82
83
84
85
            },
            'rec': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
                },
                'en': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/en_dict.txt'
                },
andyjpaddle's avatar
andyjpaddle committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
                'french': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_ppocr_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/latin_dict.txt'
                },
                'german': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_ppocr_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/latin_dict.txt'
                },
                'korean': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/korean_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/korean_dict.txt'
                },
                'japan': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/japan_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/japan_dict.txt'
                },
                'chinese_cht': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/chinese_cht_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt'
                },
                'ta': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/ta_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/ta_dict.txt'
                },
                'te': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/te_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/te_dict.txt'
                },
                'ka': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/ka_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/ka_dict.txt'
                },
                'latin': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_ppocr_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/latin_dict.txt'
                },
                'arabic': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/arabic_ppocr_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/arabic_dict.txt'
                },
                'cyrillic': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/cyrillic_ppocr_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/cyrillic_dict.txt'
                },
                'devanagari': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/devanagari_ppocr_PP-OCRv3_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/devanagari_dict.txt'
                },
146
147
148
149
150
151
152
153
            },
            'cls': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar',
                }
            },
        },
zhoujun's avatar
zhoujun committed
154
155
156
157
158
159
        'PP-OCRv2': {
            'det': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar',
                },
WenmuZhou's avatar
WenmuZhou committed
160
            },
zhoujun's avatar
zhoujun committed
161
162
163
164
165
166
            'rec': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar',
                    'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
                }
WenmuZhou's avatar
WenmuZhou committed
167
            }
WenmuZhou's avatar
WenmuZhou committed
168
        },
169
        'PP-OCR': {
zhoujun's avatar
zhoujun committed
170
171
172
173
174
175
176
177
178
179
180
181
182
            'det': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar',
                },
                'en': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar',
                },
                'structure': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar'
                }
WenmuZhou's avatar
WenmuZhou committed
183
            },
zhoujun's avatar
zhoujun committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
            'rec': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
                },
                'en': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/en_dict.txt'
                },
                'french': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/french_dict.txt'
                },
                'german': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/german_dict.txt'
                },
                'korean': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/korean_dict.txt'
                },
                'japan': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/japan_dict.txt'
                },
                'chinese_cht': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt'
                },
                'ta': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/ta_dict.txt'
                },
                'te': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/te_dict.txt'
                },
                'ka': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/ka_dict.txt'
                },
                'latin': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/latin_dict.txt'
                },
                'arabic': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/arabic_dict.txt'
                },
                'cyrillic': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/cyrillic_dict.txt'
                },
                'devanagari': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar',
                    'dict_path': './ppocr/utils/dict/devanagari_dict.txt'
                },
                'structure': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar',
                    'dict_path': 'ppocr/utils/dict/table_dict.txt'
                }
WenmuZhou's avatar
WenmuZhou committed
260
            },
zhoujun's avatar
zhoujun committed
261
262
263
264
265
            'cls': {
                'ch': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar',
                }
WenmuZhou's avatar
WenmuZhou committed
266
            },
zhoujun's avatar
zhoujun committed
267
268
269
        }
    },
    'STRUCTURE': {
270
        'PP-STRUCTURE': {
zhoujun's avatar
zhoujun committed
271
272
273
274
275
276
            'table': {
                'en': {
                    'url':
                    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar',
                    'dict_path': 'ppocr/utils/dict/table_structure_dict.txt'
                }
WenmuZhou's avatar
WenmuZhou committed
277
            }
WenmuZhou's avatar
WenmuZhou committed
278
        }
279
    }
WenmuZhou's avatar
WenmuZhou committed
280
281
282
}


WenmuZhou's avatar
WenmuZhou committed
283
def parse_args(mMain=True):
WenmuZhou's avatar
WenmuZhou committed
284
    import argparse
WenmuZhou's avatar
WenmuZhou committed
285
286
287
288
289
    parser = init_args()
    parser.add_help = mMain
    parser.add_argument("--lang", type=str, default='ch')
    parser.add_argument("--det", type=str2bool, default=True)
    parser.add_argument("--rec", type=str2bool, default=True)
290
    parser.add_argument("--type", type=str, default='ocr')
zhoujun's avatar
zhoujun committed
291
292
293
    parser.add_argument(
        "--ocr_version",
        type=str,
WenmuZhou's avatar
add hit  
WenmuZhou committed
294
        choices=SUPPORT_OCR_MODEL_VERSION,
295
        default='PP-OCRv3',
zhoujun's avatar
zhoujun committed
296
        help='OCR Model version, the current model support list is as follows: '
297
298
299
        '1. PP-OCRv3 Support Chinese and English detection and recognition model, and direction classifier model'
        '2. PP-OCRv2 Support Chinese detection and recognition model. '
        '3. PP-OCR support Chinese detection, recognition and direction classifier and multilingual recognition model.'
zhoujun's avatar
zhoujun committed
300
301
302
303
    )
    parser.add_argument(
        "--structure_version",
        type=str,
WenmuZhou's avatar
add hit  
WenmuZhou committed
304
        choices=SUPPORT_STRUCTURE_MODEL_VERSION,
305
        default='PP-STRUCTURE',
zhoujun's avatar
zhoujun committed
306
307
        help='Model version, the current model support list is as follows:'
        ' 1. STRUCTURE Support en table structure model.')
WenmuZhou's avatar
WenmuZhou committed
308
309

    for action in parser._actions:
310
        if action.dest in ['rec_char_dict_path', 'table_char_dict_path']:
WenmuZhou's avatar
WenmuZhou committed
311
            action.default = None
WenmuZhou's avatar
WenmuZhou committed
312
    if mMain:
WenmuZhou's avatar
WenmuZhou committed
313
        return parser.parse_args()
WenmuZhou's avatar
WenmuZhou committed
314
    else:
315
        inference_args_dict = {}
WenmuZhou's avatar
WenmuZhou committed
316
317
        for action in parser._actions:
            inference_args_dict[action.dest] = action.default
318
        return argparse.Namespace(**inference_args_dict)
WenmuZhou's avatar
WenmuZhou committed
319
320


WenmuZhou's avatar
WenmuZhou committed
321
322
def parse_lang(lang):
    latin_lang = [
WenmuZhou's avatar
WenmuZhou committed
323
324
325
        'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga', 'hr',
        'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi', 'ms', 'mt', 'nl',
        'no', 'oc', 'pi', 'pl', 'pt', 'ro', 'rs_latin', 'sk', 'sl', 'sq', 'sv',
andyjpaddle's avatar
andyjpaddle committed
326
        'sw', 'tl', 'tr', 'uz', 'vi', 'french', 'german'
WenmuZhou's avatar
WenmuZhou committed
327
328
329
    ]
    arabic_lang = ['ar', 'fa', 'ug', 'ur']
    cyrillic_lang = [
WenmuZhou's avatar
WenmuZhou committed
330
331
        'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',
        'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
WenmuZhou's avatar
WenmuZhou committed
332
333
    ]
    devanagari_lang = [
WenmuZhou's avatar
WenmuZhou committed
334
335
        'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',
        'sa', 'bgc'
WenmuZhou's avatar
WenmuZhou committed
336
337
338
339
340
341
342
343
344
    ]
    if lang in latin_lang:
        lang = "latin"
    elif lang in arabic_lang:
        lang = "arabic"
    elif lang in cyrillic_lang:
        lang = "cyrillic"
    elif lang in devanagari_lang:
        lang = "devanagari"
zhoujun's avatar
zhoujun committed
345
    assert lang in MODEL_URLS['OCR'][DEFAULT_OCR_MODEL_VERSION][
WenmuZhou's avatar
WenmuZhou committed
346
        'rec'], 'param lang must in {}, but got {}'.format(
zhoujun's avatar
zhoujun committed
347
            MODEL_URLS['OCR'][DEFAULT_OCR_MODEL_VERSION]['rec'].keys(), lang)
WenmuZhou's avatar
WenmuZhou committed
348
349
350
351
    if lang == "ch":
        det_lang = "ch"
    elif lang == 'structure':
        det_lang = 'structure'
andyjpaddle's avatar
andyjpaddle committed
352
353
    elif lang == "ml":
        det_lang = "ml"
WenmuZhou's avatar
WenmuZhou committed
354
355
356
357
358
    else:
        det_lang = "en"
    return lang, det_lang


zhoujun's avatar
zhoujun committed
359
360
361
362
363
364
365
def get_model_config(type, version, model_type, lang):
    if type == 'OCR':
        DEFAULT_MODEL_VERSION = DEFAULT_OCR_MODEL_VERSION
    elif type == 'STRUCTURE':
        DEFAULT_MODEL_VERSION = DEFAULT_STRUCTURE_MODEL_VERSION
    else:
        raise NotImplementedError
WenmuZhou's avatar
WenmuZhou committed
366

zhoujun's avatar
zhoujun committed
367
368
    model_urls = MODEL_URLS[type]
    if version not in model_urls:
WenmuZhou's avatar
WenmuZhou committed
369
        version = DEFAULT_MODEL_VERSION
zhoujun's avatar
zhoujun committed
370
371
    if model_type not in model_urls[version]:
        if model_type in model_urls[DEFAULT_MODEL_VERSION]:
WenmuZhou's avatar
WenmuZhou committed
372
373
374
            version = DEFAULT_MODEL_VERSION
        else:
            logger.error('{} models is not support, we only support {}'.format(
zhoujun's avatar
zhoujun committed
375
                model_type, model_urls[DEFAULT_MODEL_VERSION].keys()))
WenmuZhou's avatar
WenmuZhou committed
376
            sys.exit(-1)
WenmuZhou's avatar
add hit  
WenmuZhou committed
377

zhoujun's avatar
zhoujun committed
378
379
    if lang not in model_urls[version][model_type]:
        if lang in model_urls[DEFAULT_MODEL_VERSION][model_type]:
WenmuZhou's avatar
WenmuZhou committed
380
381
382
383
            version = DEFAULT_MODEL_VERSION
        else:
            logger.error(
                'lang {} is not support, we only support {} for {} models'.
zhoujun's avatar
zhoujun committed
384
                format(lang, model_urls[DEFAULT_MODEL_VERSION][model_type].keys(
WenmuZhou's avatar
WenmuZhou committed
385
386
                ), model_type))
            sys.exit(-1)
zhoujun's avatar
zhoujun committed
387
    return model_urls[version][model_type][lang]
WenmuZhou's avatar
WenmuZhou committed
388
389


WenmuZhou's avatar
WenmuZhou committed
390
class PaddleOCR(predict_system.TextSystem):
391
    def __init__(self, **kwargs):
WenmuZhou's avatar
WenmuZhou committed
392
393
394
395
396
        """
        paddleocr package
        args:
            **kwargs: other params show in paddleocr --help
        """
WenmuZhou's avatar
WenmuZhou committed
397
398
        params = parse_args(mMain=False)
        params.__dict__.update(**kwargs)
WenmuZhou's avatar
add hit  
WenmuZhou committed
399
400
        assert params.ocr_version in SUPPORT_OCR_MODEL_VERSION, "ocr_version must in {}, but get {}".format(
            SUPPORT_OCR_MODEL_VERSION, params.ocr_version)
zhoujun's avatar
zhoujun committed
401
402
        params.use_gpu = check_gpu(params.use_gpu)

WenmuZhou's avatar
WenmuZhou committed
403
404
        if not params.show_log:
            logger.setLevel(logging.INFO)
WenmuZhou's avatar
WenmuZhou committed
405
        self.use_angle_cls = params.use_angle_cls
WenmuZhou's avatar
WenmuZhou committed
406
        lang, det_lang = parse_lang(params.lang)
WenmuZhou's avatar
WenmuZhou committed
407

408
        # init model dir
zhoujun's avatar
zhoujun committed
409
410
        det_model_config = get_model_config('OCR', params.ocr_version, 'det',
                                            det_lang)
WenmuZhou's avatar
WenmuZhou committed
411
412
        params.det_model_dir, det_url = confirm_model_dir_url(
            params.det_model_dir,
WenmuZhou's avatar
WenmuZhou committed
413
            os.path.join(BASE_DIR, 'whl', 'det', det_lang),
WenmuZhou's avatar
WenmuZhou committed
414
            det_model_config['url'])
zhoujun's avatar
zhoujun committed
415
416
        rec_model_config = get_model_config('OCR', params.ocr_version, 'rec',
                                            lang)
WenmuZhou's avatar
WenmuZhou committed
417
418
        params.rec_model_dir, rec_url = confirm_model_dir_url(
            params.rec_model_dir,
WenmuZhou's avatar
WenmuZhou committed
419
            os.path.join(BASE_DIR, 'whl', 'rec', lang), rec_model_config['url'])
zhoujun's avatar
zhoujun committed
420
421
        cls_model_config = get_model_config('OCR', params.ocr_version, 'cls',
                                            'ch')
WenmuZhou's avatar
WenmuZhou committed
422
423
        params.cls_model_dir, cls_url = confirm_model_dir_url(
            params.cls_model_dir,
WenmuZhou's avatar
WenmuZhou committed
424
            os.path.join(BASE_DIR, 'whl', 'cls'), cls_model_config['url'])
WenmuZhou's avatar
WenmuZhou committed
425
        # download model
426
427
428
        maybe_download(params.det_model_dir, det_url)
        maybe_download(params.rec_model_dir, rec_url)
        maybe_download(params.cls_model_dir, cls_url)
WenmuZhou's avatar
WenmuZhou committed
429

WenmuZhou's avatar
WenmuZhou committed
430
        if params.det_algorithm not in SUPPORT_DET_MODEL:
WenmuZhou's avatar
WenmuZhou committed
431
432
            logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
            sys.exit(0)
WenmuZhou's avatar
WenmuZhou committed
433
        if params.rec_algorithm not in SUPPORT_REC_MODEL:
WenmuZhou's avatar
WenmuZhou committed
434
435
            logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
            sys.exit(0)
WenmuZhou's avatar
WenmuZhou committed
436
437

        if params.rec_char_dict_path is None:
WenmuZhou's avatar
WenmuZhou committed
438
439
            params.rec_char_dict_path = str(
                Path(__file__).parent / rec_model_config['dict_path'])
WenmuZhou's avatar
WenmuZhou committed
440

WenmuZhou's avatar
WenmuZhou committed
441
        logger.debug(params)
WenmuZhou's avatar
WenmuZhou committed
442
        # init det_model and rec_model
WenmuZhou's avatar
WenmuZhou committed
443
        super().__init__(params)
WenmuZhou's avatar
WenmuZhou committed
444

445
    def ocr(self, img, det=True, rec=True, cls=True):
WenmuZhou's avatar
WenmuZhou committed
446
447
448
449
        """
        ocr with paddleocr
        args:
            img: img for ocr, support ndarray, img_path and list or ndarray
fanruinet's avatar
fanruinet committed
450
451
452
            det: use text detection or not. If false, only rec will be exec. Default is True
            rec: use text recognition or not. If false, only det will be exec. Default is True
            cls: use angle classifier or not. Default is True. If true, the text with rotation of 180 degrees can be recognized. If no text is rotated by 180 degrees, use cls=False to get better performance. Text with rotation of 90 or 270 degrees can be recognized even if cls=False.
WenmuZhou's avatar
WenmuZhou committed
453
454
        """
        assert isinstance(img, (np.ndarray, list, str))
WenmuZhou's avatar
WenmuZhou committed
455
456
457
        if isinstance(img, list) and det == True:
            logger.error('When input a list of images, det must be false')
            exit(0)
458
        if cls == True and self.use_angle_cls == False:
WenmuZhou's avatar
WenmuZhou committed
459
460
461
            logger.warning(
                'Since the angle classifier is not initialized, the angle classifier will not be uesd during the forward process'
            )
WenmuZhou's avatar
WenmuZhou committed
462

WenmuZhou's avatar
WenmuZhou committed
463
        if isinstance(img, str):
WenmuZhou's avatar
WenmuZhou committed
464
465
466
467
            # download net image
            if img.startswith('http'):
                download_with_progressbar(img, 'tmp.jpg')
                img = 'tmp.jpg'
WenmuZhou's avatar
WenmuZhou committed
468
469
470
            image_file = img
            img, flag = check_and_read_gif(image_file)
            if not flag:
471
472
473
                with open(image_file, 'rb') as f:
                    np_arr = np.frombuffer(f.read(), dtype=np.uint8)
                    img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
WenmuZhou's avatar
WenmuZhou committed
474
475
476
            if img is None:
                logger.error("error in loading image:{}".format(image_file))
                return None
WenmuZhou's avatar
WenmuZhou committed
477
478
        if isinstance(img, np.ndarray) and len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
WenmuZhou's avatar
WenmuZhou committed
479
        if det and rec:
480
            dt_boxes, rec_res = self.__call__(img, cls)
WenmuZhou's avatar
WenmuZhou committed
481
482
483
484
485
486
487
488
489
            return [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
        elif det and not rec:
            dt_boxes, elapse = self.text_detector(img)
            if dt_boxes is None:
                return None
            return [box.tolist() for box in dt_boxes]
        else:
            if not isinstance(img, list):
                img = [img]
490
            if self.use_angle_cls and cls:
WenmuZhou's avatar
WenmuZhou committed
491
492
493
                img, cls_res, elapse = self.text_classifier(img)
                if not rec:
                    return cls_res
WenmuZhou's avatar
WenmuZhou committed
494
495
            rec_res, elapse = self.text_recognizer(img)
            return rec_res
496
497


498
class PPStructure(StructureSystem):
499
500
501
    def __init__(self, **kwargs):
        params = parse_args(mMain=False)
        params.__dict__.update(**kwargs)
WenmuZhou's avatar
add hit  
WenmuZhou committed
502
        assert params.structure_version in SUPPORT_STRUCTURE_MODEL_VERSION, "structure_version must in {}, but get {}".format(
WenmuZhou's avatar
add hit  
WenmuZhou committed
503
            SUPPORT_STRUCTURE_MODEL_VERSION, params.structure_version)
zhoujun's avatar
zhoujun committed
504
505
        params.use_gpu = check_gpu(params.use_gpu)

506
507
        if not params.show_log:
            logger.setLevel(logging.INFO)
WenmuZhou's avatar
WenmuZhou committed
508
509
        lang, det_lang = parse_lang(params.lang)

510
        # init model dir
zhoujun's avatar
zhoujun committed
511
512
        det_model_config = get_model_config('OCR', params.ocr_version, 'det',
                                            det_lang)
WenmuZhou's avatar
WenmuZhou committed
513
514
        params.det_model_dir, det_url = confirm_model_dir_url(
            params.det_model_dir,
WenmuZhou's avatar
WenmuZhou committed
515
            os.path.join(BASE_DIR, 'whl', 'det', det_lang),
WenmuZhou's avatar
WenmuZhou committed
516
            det_model_config['url'])
zhoujun's avatar
zhoujun committed
517
518
        rec_model_config = get_model_config('OCR', params.ocr_version, 'rec',
                                            lang)
WenmuZhou's avatar
WenmuZhou committed
519
520
        params.rec_model_dir, rec_url = confirm_model_dir_url(
            params.rec_model_dir,
WenmuZhou's avatar
WenmuZhou committed
521
            os.path.join(BASE_DIR, 'whl', 'rec', lang), rec_model_config['url'])
zhoujun's avatar
zhoujun committed
522
523
        table_model_config = get_model_config(
            'STRUCTURE', params.structure_version, 'table', 'en')
WenmuZhou's avatar
WenmuZhou committed
524
525
        params.table_model_dir, table_url = confirm_model_dir_url(
            params.table_model_dir,
WenmuZhou's avatar
WenmuZhou committed
526
            os.path.join(BASE_DIR, 'whl', 'table'), table_model_config['url'])
527
528
529
530
531
532
        # download model
        maybe_download(params.det_model_dir, det_url)
        maybe_download(params.rec_model_dir, rec_url)
        maybe_download(params.table_model_dir, table_url)

        if params.rec_char_dict_path is None:
WenmuZhou's avatar
WenmuZhou committed
533
534
            params.rec_char_dict_path = str(
                Path(__file__).parent / rec_model_config['dict_path'])
535
        if params.table_char_dict_path is None:
WenmuZhou's avatar
WenmuZhou committed
536
537
            params.table_char_dict_path = str(
                Path(__file__).parent / table_model_config['dict_path'])
538

WenmuZhou's avatar
WenmuZhou committed
539
        logger.debug(params)
540
541
        super().__init__(params)

542
    def __call__(self, img, return_ocr_result_in_table=False):
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
        if isinstance(img, str):
            # download net image
            if img.startswith('http'):
                download_with_progressbar(img, 'tmp.jpg')
                img = 'tmp.jpg'
            image_file = img
            img, flag = check_and_read_gif(image_file)
            if not flag:
                with open(image_file, 'rb') as f:
                    np_arr = np.frombuffer(f.read(), dtype=np.uint8)
                    img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            if img is None:
                logger.error("error in loading image:{}".format(image_file))
                return None
        if isinstance(img, np.ndarray) and len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

560
        res = super().__call__(img, return_ocr_result_in_table)
561
562
563
        return res


564
def main():
WenmuZhou's avatar
WenmuZhou committed
565
    # for cmd
WenmuZhou's avatar
WenmuZhou committed
566
    args = parse_args(mMain=True)
WenmuZhou's avatar
WenmuZhou committed
567
    image_dir = args.image_dir
568
    if is_link(image_dir):
WenmuZhou's avatar
WenmuZhou committed
569
570
571
572
        download_with_progressbar(image_dir, 'tmp.jpg')
        image_file_list = ['tmp.jpg']
    else:
        image_file_list = get_image_file_list(args.image_dir)
573
574
575
    if len(image_file_list) == 0:
        logger.error('no images find in {}'.format(args.image_dir))
        return
WenmuZhou's avatar
WenmuZhou committed
576
    if args.type == 'ocr':
577
        engine = PaddleOCR(**(args.__dict__))
WenmuZhou's avatar
WenmuZhou committed
578
    elif args.type == 'structure':
579
580
581
        engine = PPStructure(**(args.__dict__))
    else:
        raise NotImplementedError
WenmuZhou's avatar
WenmuZhou committed
582

583
    for img_path in image_file_list:
584
        img_name = os.path.basename(img_path).split('.')[0]
WenmuZhou's avatar
WenmuZhou committed
585
        logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
586
587
        if args.type == 'ocr':
            result = engine.ocr(img_path,
WenmuZhou's avatar
WenmuZhou committed
588
589
590
                                det=args.det,
                                rec=args.rec,
                                cls=args.use_angle_cls)
591
592
593
594
595
            if result is not None:
                for line in result:
                    logger.info(line)
        elif args.type == 'structure':
            result = engine(img_path)
WenmuZhou's avatar
WenmuZhou committed
596
597
            save_structure_res(result, args.output, img_name)

598
            for item in result:
WenmuZhou's avatar
WenmuZhou committed
599
                item.pop('img')
WenmuZhou's avatar
WenmuZhou committed
600
                logger.info(item)