paddleocr.py 17.8 KB
Newer Older
WenmuZhou's avatar
WenmuZhou committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 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

__dir__ = os.path.dirname(__file__)
sys.path.append(os.path.join(__dir__, ''))

import cv2
WenmuZhou's avatar
WenmuZhou committed
22
import logging
WenmuZhou's avatar
WenmuZhou committed
23
24
25
26
import numpy as np
from pathlib import Path

from tools.infer import predict_system
WenmuZhou's avatar
WenmuZhou committed
27
from ppocr.utils.logging import get_logger
WenmuZhou's avatar
WenmuZhou committed
28

WenmuZhou's avatar
WenmuZhou committed
29
logger = get_logger()
30
from ppocr.utils.utility import check_and_read_gif, get_image_file_list
31
from ppocr.utils.network import maybe_download, download_with_progressbar, is_link, confirm_model_dir_url
32
33
34
from tools.infer.utility import draw_ocr, str2bool
from ppstructure.utility import init_args, draw_structure_result
from ppstructure.predict_system import OCRSystem, save_structure_res
WenmuZhou's avatar
WenmuZhou committed
35

WenmuZhou's avatar
WenmuZhou committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
__all__ = [
    'PaddleOCR', 'PPStructure', 'draw_ocr', 'draw_structure_result',
    'save_structure_res', 'download_with_progressbar'
]

SUPPORT_DET_MODEL = ['DB']
VERSION = '2.2.1'
SUPPORT_REC_MODEL = ['CRNN']
BASE_DIR = os.path.expanduser("~/.paddleocr/")

DEFAULT_MODEL_VERSION = '2.0'
MODEL_URLS = {
    '2.1': {
        'det': {
            'ch': {
                'url':
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
52
                'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
53
54
55
56
57
            },
        },
        'rec': {
            'ch': {
                'url':
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
58
                'https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
59
60
61
                'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
            }
        }
tink2123's avatar
tink2123 committed
62
    },
WenmuZhou's avatar
WenmuZhou committed
63
64
65
66
67
68
69
70
71
72
73
74
75
76
    '2.0': {
        '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
77
        },
WenmuZhou's avatar
WenmuZhou committed
78
79
80
81
82
83
84
85
        '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':
WenmuZhou's avatar
WenmuZhou committed
86
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
87
88
89
90
                'dict_path': './ppocr/utils/en_dict.txt'
            },
            'french': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
91
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
92
93
94
95
                'dict_path': './ppocr/utils/dict/french_dict.txt'
            },
            'german': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
96
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
97
98
99
100
                'dict_path': './ppocr/utils/dict/german_dict.txt'
            },
            'korean': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
101
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
102
103
104
105
                'dict_path': './ppocr/utils/dict/korean_dict.txt'
            },
            'japan': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
106
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
107
108
109
110
                'dict_path': './ppocr/utils/dict/japan_dict.txt'
            },
            'chinese_cht': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
111
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
112
113
114
115
                'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt'
            },
            'ta': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
116
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
117
118
119
120
                'dict_path': './ppocr/utils/dict/ta_dict.txt'
            },
            'te': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
121
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
122
123
124
125
                'dict_path': './ppocr/utils/dict/te_dict.txt'
            },
            'ka': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
126
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
127
128
129
130
                'dict_path': './ppocr/utils/dict/ka_dict.txt'
            },
            'latin': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
131
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
132
133
134
135
                'dict_path': './ppocr/utils/dict/latin_dict.txt'
            },
            'arabic': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
136
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
137
138
139
140
                'dict_path': './ppocr/utils/dict/arabic_dict.txt'
            },
            'cyrillic': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
141
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
142
143
144
145
                'dict_path': './ppocr/utils/dict/cyrillic_dict.txt'
            },
            'devanagari': {
                'url':
WenmuZhou's avatar
WenmuZhou committed
146
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
147
148
149
150
151
152
153
154
155
156
157
158
159
                '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'
            }
        },
        'cls': {
            'ch': {
                'url':
                'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar',
            }
160
        },
WenmuZhou's avatar
WenmuZhou committed
161
162
163
164
165
166
        '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
167
        }
168
    }
WenmuZhou's avatar
WenmuZhou committed
169
170
171
}


WenmuZhou's avatar
WenmuZhou committed
172
def parse_args(mMain=True):
WenmuZhou's avatar
WenmuZhou committed
173
    import argparse
WenmuZhou's avatar
WenmuZhou committed
174
175
176
177
178
    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)
179
    parser.add_argument("--type", type=str, default='ocr')
WenmuZhou's avatar
WenmuZhou committed
180
    parser.add_argument("--version", type=str, default='2.1')
WenmuZhou's avatar
WenmuZhou committed
181
182

    for action in parser._actions:
183
        if action.dest in ['rec_char_dict_path', 'table_char_dict_path']:
WenmuZhou's avatar
WenmuZhou committed
184
            action.default = None
WenmuZhou's avatar
WenmuZhou committed
185
    if mMain:
WenmuZhou's avatar
WenmuZhou committed
186
        return parser.parse_args()
WenmuZhou's avatar
WenmuZhou committed
187
    else:
188
        inference_args_dict = {}
WenmuZhou's avatar
WenmuZhou committed
189
190
        for action in parser._actions:
            inference_args_dict[action.dest] = action.default
191
        return argparse.Namespace(**inference_args_dict)
WenmuZhou's avatar
WenmuZhou committed
192
193


WenmuZhou's avatar
WenmuZhou committed
194
195
def parse_lang(lang):
    latin_lang = [
WenmuZhou's avatar
WenmuZhou committed
196
197
198
199
        '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',
        'sw', 'tl', 'tr', 'uz', 'vi'
WenmuZhou's avatar
WenmuZhou committed
200
201
202
    ]
    arabic_lang = ['ar', 'fa', 'ug', 'ur']
    cyrillic_lang = [
WenmuZhou's avatar
WenmuZhou committed
203
204
        'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',
        'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
WenmuZhou's avatar
WenmuZhou committed
205
206
    ]
    devanagari_lang = [
WenmuZhou's avatar
WenmuZhou committed
207
208
        'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',
        'sa', 'bgc'
WenmuZhou's avatar
WenmuZhou committed
209
210
211
212
213
214
215
216
217
    ]
    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"
WenmuZhou's avatar
WenmuZhou committed
218
    assert lang in MODEL_URLS[DEFAULT_MODEL_VERSION][
WenmuZhou's avatar
WenmuZhou committed
219
        'rec'], 'param lang must in {}, but got {}'.format(
WenmuZhou's avatar
WenmuZhou committed
220
            MODEL_URLS[DEFAULT_MODEL_VERSION]['rec'].keys(), lang)
WenmuZhou's avatar
WenmuZhou committed
221
222
223
224
225
226
227
228
229
    if lang == "ch":
        det_lang = "ch"
    elif lang == 'structure':
        det_lang = 'structure'
    else:
        det_lang = "en"
    return lang, det_lang


WenmuZhou's avatar
WenmuZhou committed
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
def get_model_config(version, model_type, lang):
    if version not in MODEL_URLS:
        logger.warning('version {} not in {}, use version {} instead'.format(
            version, MODEL_URLS.keys(), DEFAULT_MODEL_VERSION))
        version = DEFAULT_MODEL_VERSION
    if model_type not in MODEL_URLS[version]:
        if model_type in MODEL_URLS[DEFAULT_MODEL_VERSION]:
            logger.warning(
                'version {} not support {} models, use version {} instead'.
                format(version, model_type, DEFAULT_MODEL_VERSION))
            version = DEFAULT_MODEL_VERSION
        else:
            logger.error('{} models is not support, we only support {}'.format(
                model_type, MODEL_URLS[DEFAULT_MODEL_VERSION].keys()))
            sys.exit(-1)
    if lang not in MODEL_URLS[version][model_type]:
        if lang in MODEL_URLS[DEFAULT_MODEL_VERSION][model_type]:
            logger.warning('lang {} is not support in {}, use {} instead'.
                           format(lang, version, DEFAULT_MODEL_VERSION))
            version = DEFAULT_MODEL_VERSION
        else:
            logger.error(
                'lang {} is not support, we only support {} for {} models'.
                format(lang, MODEL_URLS[DEFAULT_MODEL_VERSION][model_type].keys(
                ), model_type))
            sys.exit(-1)
    return MODEL_URLS[version][model_type][lang]


WenmuZhou's avatar
WenmuZhou committed
259
class PaddleOCR(predict_system.TextSystem):
260
    def __init__(self, **kwargs):
WenmuZhou's avatar
WenmuZhou committed
261
262
263
264
265
        """
        paddleocr package
        args:
            **kwargs: other params show in paddleocr --help
        """
WenmuZhou's avatar
WenmuZhou committed
266
267
        params = parse_args(mMain=False)
        params.__dict__.update(**kwargs)
WenmuZhou's avatar
WenmuZhou committed
268
269
        if not params.show_log:
            logger.setLevel(logging.INFO)
WenmuZhou's avatar
WenmuZhou committed
270
        self.use_angle_cls = params.use_angle_cls
WenmuZhou's avatar
WenmuZhou committed
271
        lang, det_lang = parse_lang(params.lang)
WenmuZhou's avatar
WenmuZhou committed
272

273
        # init model dir
WenmuZhou's avatar
WenmuZhou committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
        det_model_config = get_model_config(params.version, 'det', det_lang)
        params.det_model_dir, det_url = confirm_model_dir_url(
            params.det_model_dir,
            os.path.join(BASE_DIR, VERSION, 'ocr', 'det', det_lang),
            det_model_config['url'])
        rec_model_config = get_model_config(params.version, 'rec', lang)
        params.rec_model_dir, rec_url = confirm_model_dir_url(
            params.rec_model_dir,
            os.path.join(BASE_DIR, VERSION, 'ocr', 'rec', lang),
            rec_model_config['url'])
        cls_model_config = get_model_config(params.version, 'cls', 'ch')
        params.cls_model_dir, cls_url = confirm_model_dir_url(
            params.cls_model_dir,
            os.path.join(BASE_DIR, VERSION, 'ocr', 'cls'),
            cls_model_config['url'])
WenmuZhou's avatar
WenmuZhou committed
289
        # download model
290
291
292
        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
293

WenmuZhou's avatar
WenmuZhou committed
294
        if params.det_algorithm not in SUPPORT_DET_MODEL:
WenmuZhou's avatar
WenmuZhou committed
295
296
            logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
            sys.exit(0)
WenmuZhou's avatar
WenmuZhou committed
297
        if params.rec_algorithm not in SUPPORT_REC_MODEL:
WenmuZhou's avatar
WenmuZhou committed
298
299
            logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
            sys.exit(0)
WenmuZhou's avatar
WenmuZhou committed
300
301

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

WenmuZhou's avatar
WenmuZhou committed
305
        print(params)
WenmuZhou's avatar
WenmuZhou committed
306
        # init det_model and rec_model
WenmuZhou's avatar
WenmuZhou committed
307
        super().__init__(params)
WenmuZhou's avatar
WenmuZhou committed
308

309
    def ocr(self, img, det=True, rec=True, cls=True):
WenmuZhou's avatar
WenmuZhou committed
310
311
312
313
314
315
316
317
        """
        ocr with paddleocr
        args:
            img: img for ocr, support ndarray, img_path and list or ndarray
            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
        """
        assert isinstance(img, (np.ndarray, list, str))
WenmuZhou's avatar
WenmuZhou committed
318
319
320
        if isinstance(img, list) and det == True:
            logger.error('When input a list of images, det must be false')
            exit(0)
321
        if cls == True and self.use_angle_cls == False:
WenmuZhou's avatar
WenmuZhou committed
322
323
324
            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
325

WenmuZhou's avatar
WenmuZhou committed
326
        if isinstance(img, str):
WenmuZhou's avatar
WenmuZhou committed
327
328
329
330
            # download net image
            if img.startswith('http'):
                download_with_progressbar(img, 'tmp.jpg')
                img = 'tmp.jpg'
WenmuZhou's avatar
WenmuZhou committed
331
332
333
            image_file = img
            img, flag = check_and_read_gif(image_file)
            if not flag:
334
335
336
                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
337
338
339
            if img is None:
                logger.error("error in loading image:{}".format(image_file))
                return None
WenmuZhou's avatar
WenmuZhou committed
340
341
        if isinstance(img, np.ndarray) and len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
WenmuZhou's avatar
WenmuZhou committed
342
        if det and rec:
343
            dt_boxes, rec_res = self.__call__(img, cls)
WenmuZhou's avatar
WenmuZhou committed
344
345
346
347
348
349
350
351
352
            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]
353
            if self.use_angle_cls and cls:
WenmuZhou's avatar
WenmuZhou committed
354
355
356
                img, cls_res, elapse = self.text_classifier(img)
                if not rec:
                    return cls_res
WenmuZhou's avatar
WenmuZhou committed
357
358
            rec_res, elapse = self.text_recognizer(img)
            return rec_res
359
360


361
362
363
364
365
366
class PPStructure(OCRSystem):
    def __init__(self, **kwargs):
        params = parse_args(mMain=False)
        params.__dict__.update(**kwargs)
        if not params.show_log:
            logger.setLevel(logging.INFO)
WenmuZhou's avatar
WenmuZhou committed
367
368
        lang, det_lang = parse_lang(params.lang)

369
        # init model dir
WenmuZhou's avatar
WenmuZhou committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
        det_model_config = get_model_config(params.version, 'det', det_lang)
        params.det_model_dir, det_url = confirm_model_dir_url(
            params.det_model_dir,
            os.path.join(BASE_DIR, VERSION, 'ocr', 'det', det_lang),
            det_model_config['url'])
        rec_model_config = get_model_config(params.version, 'rec', lang)
        params.rec_model_dir, rec_url = confirm_model_dir_url(
            params.rec_model_dir,
            os.path.join(BASE_DIR, VERSION, 'ocr', 'rec', lang),
            rec_model_config['url'])
        table_model_config = get_model_config(params.version, 'table', 'en')
        params.table_model_dir, table_url = confirm_model_dir_url(
            params.table_model_dir,
            os.path.join(BASE_DIR, VERSION, 'ocr', 'table'),
            table_model_config['url'])
385
386
387
388
389
390
        # 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
391
392
            params.rec_char_dict_path = str(
                Path(__file__).parent / rec_model_config['dict_path'])
393
        if params.table_char_dict_path is None:
WenmuZhou's avatar
WenmuZhou committed
394
395
            params.table_char_dict_path = str(
                Path(__file__).parent / table_model_config['dict_path'])
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421

        print(params)
        super().__init__(params)

    def __call__(self, img):
        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)

        res = super().__call__(img)
        return res


422
def main():
WenmuZhou's avatar
WenmuZhou committed
423
    # for cmd
WenmuZhou's avatar
WenmuZhou committed
424
    args = parse_args(mMain=True)
WenmuZhou's avatar
WenmuZhou committed
425
    image_dir = args.image_dir
426
    if is_link(image_dir):
WenmuZhou's avatar
WenmuZhou committed
427
428
429
430
        download_with_progressbar(image_dir, 'tmp.jpg')
        image_file_list = ['tmp.jpg']
    else:
        image_file_list = get_image_file_list(args.image_dir)
431
432
433
    if len(image_file_list) == 0:
        logger.error('no images find in {}'.format(args.image_dir))
        return
WenmuZhou's avatar
WenmuZhou committed
434
    if args.type == 'ocr':
435
        engine = PaddleOCR(**(args.__dict__))
WenmuZhou's avatar
WenmuZhou committed
436
    elif args.type == 'structure':
437
438
439
        engine = PPStructure(**(args.__dict__))
    else:
        raise NotImplementedError
WenmuZhou's avatar
WenmuZhou committed
440

441
    for img_path in image_file_list:
442
        img_name = os.path.basename(img_path).split('.')[0]
WenmuZhou's avatar
WenmuZhou committed
443
        logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
444
445
        if args.type == 'ocr':
            result = engine.ocr(img_path,
WenmuZhou's avatar
WenmuZhou committed
446
447
448
                                det=args.det,
                                rec=args.rec,
                                cls=args.use_angle_cls)
449
450
451
452
453
            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
454
455
            save_structure_res(result, args.output, img_name)

456
            for item in result:
WenmuZhou's avatar
WenmuZhou committed
457
                item.pop('img')
MissPenguin's avatar
MissPenguin committed
458
                logger.info(item)