paddleocr.py 13.5 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
22
23
24
25
26
27
28
# 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
import numpy as np
from pathlib import Path
import tarfile
import requests
from tqdm import tqdm

from tools.infer import predict_system
WenmuZhou's avatar
WenmuZhou committed
29
from ppocr.utils.logging import get_logger
WenmuZhou's avatar
WenmuZhou committed
30

WenmuZhou's avatar
WenmuZhou committed
31
logger = get_logger()
32
from ppocr.utils.utility import check_and_read_gif, get_image_file_list
33
from tools.infer.utility import draw_ocr, inference_args_list, str2bool, parse_args
WenmuZhou's avatar
WenmuZhou committed
34
35
36

__all__ = ['PaddleOCR']

WenmuZhou's avatar
WenmuZhou committed
37
model_urls = {
tink2123's avatar
tink2123 committed
38
39
40
41
42
43
    'det': {
        'ch':
        'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar',
        'en':
        'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar'
    },
WenmuZhou's avatar
WenmuZhou committed
44
45
46
    'rec': {
        'ch': {
            'url':
WenmuZhou's avatar
WenmuZhou committed
47
            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
48
49
50
51
            'dict_path': './ppocr/utils/ppocr_keys_v1.txt'
        },
        'en': {
            'url':
WenmuZhou's avatar
WenmuZhou committed
52
            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar',
tink2123's avatar
tink2123 committed
53
            'dict_path': './ppocr/utils/en_dict.txt'
WenmuZhou's avatar
WenmuZhou committed
54
55
56
        },
        'french': {
            'url':
WenmuZhou's avatar
WenmuZhou committed
57
            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
58
59
60
61
            'dict_path': './ppocr/utils/dict/french_dict.txt'
        },
        'german': {
            'url':
WenmuZhou's avatar
WenmuZhou committed
62
            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
63
64
65
66
            'dict_path': './ppocr/utils/dict/german_dict.txt'
        },
        'korean': {
            'url':
WenmuZhou's avatar
WenmuZhou committed
67
            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
68
69
70
71
            'dict_path': './ppocr/utils/dict/korean_dict.txt'
        },
        'japan': {
            'url':
WenmuZhou's avatar
WenmuZhou committed
72
            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
WenmuZhou's avatar
WenmuZhou committed
73
            'dict_path': './ppocr/utils/dict/japan_dict.txt'
tink2123's avatar
tink2123 committed
74
75
76
77
78
79
80
81
82
83
84
85
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
        },
        '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'
WenmuZhou's avatar
WenmuZhou committed
114
115
116
        }
    },
    'cls':
WenmuZhou's avatar
WenmuZhou committed
117
    'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar'
WenmuZhou's avatar
WenmuZhou committed
118
119
120
}

SUPPORT_DET_MODEL = ['DB']
tink2123's avatar
tink2123 committed
121
VERSION = '2.1'
122
123
SUPPORT_REC_MODEL = ['CRNN']
BASE_DIR = os.path.expanduser("~/.paddleocr/")
WenmuZhou's avatar
WenmuZhou committed
124
125
126
127
128
129
130
131
132
133
134
135


def download_with_progressbar(url, save_path):
    response = requests.get(url, stream=True)
    total_size_in_bytes = int(response.headers.get('content-length', 0))
    block_size = 1024  # 1 Kibibyte
    progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
    with open(save_path, 'wb') as file:
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)
    progress_bar.close()
WenmuZhou's avatar
WenmuZhou committed
136
137
    if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes:
        logger.error("Something went wrong while downloading models")
WenmuZhou's avatar
WenmuZhou committed
138
139
140
        sys.exit(0)


141
def maybe_download(model_storage_directory, url):
WenmuZhou's avatar
WenmuZhou committed
142
    # using custom model
WenmuZhou's avatar
WenmuZhou committed
143
144
145
146
147
148
149
    tar_file_name_list = [
        'inference.pdiparams', 'inference.pdiparams.info', 'inference.pdmodel'
    ]
    if not os.path.exists(
            os.path.join(model_storage_directory, 'inference.pdiparams')
    ) or not os.path.exists(
            os.path.join(model_storage_directory, 'inference.pdmodel')):
150
151
152
153
154
155
        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():
WenmuZhou's avatar
WenmuZhou committed
156
157
158
159
160
                filename = None
                for tar_file_name in tar_file_name_list:
                    if tar_file_name in member.name:
                        filename = tar_file_name
                if filename is None:
161
162
163
164
165
166
167
                    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)
WenmuZhou's avatar
WenmuZhou committed
168
169


170
def parse_args_whl(mMain=True):
WenmuZhou's avatar
WenmuZhou committed
171
    import argparse
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    extend_args_list = [
        {
            'name': 'lang',
            'type': str,
            'default': 'ch'
        },
        {
            'name': 'det',
            'type': str2bool,
            'default': True
        },
        {
            'name': 'rec',
            'type': str2bool,
            'default': True
        },
    ]
    for item in inference_args_list:
        if item['name'] == 'rec_char_dict_path':
            item['default'] = None
    inference_args_list.extend(extend_args_list)
WenmuZhou's avatar
WenmuZhou committed
193
    if mMain:
194
        return parse_args()
WenmuZhou's avatar
WenmuZhou committed
195
    else:
196
197
198
199
        inference_args_dict = {}
        for item in inference_args_list:
            inference_args_dict[item['name']] = item['default']
        return argparse.Namespace(**inference_args_dict)
WenmuZhou's avatar
WenmuZhou committed
200
201
202


class PaddleOCR(predict_system.TextSystem):
203
    def __init__(self, **kwargs):
WenmuZhou's avatar
WenmuZhou committed
204
205
206
207
208
        """
        paddleocr package
        args:
            **kwargs: other params show in paddleocr --help
        """
209
        postprocess_params = parse_args_whl(mMain=False)
210
        postprocess_params.__dict__.update(**kwargs)
WenmuZhou's avatar
WenmuZhou committed
211
212
        self.use_angle_cls = postprocess_params.use_angle_cls
        lang = postprocess_params.lang
tink2123's avatar
tink2123 committed
213
        latin_lang = [
tink2123's avatar
tink2123 committed
214
215
216
217
            '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'
tink2123's avatar
tink2123 committed
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
        ]
        arabic_lang = ['ar', 'fa', 'ug', 'ur']
        cyrillic_lang = [
            'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd',
            'ava', 'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
        ]
        devanagari_lang = [
            'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new',
            'gom', 'sa', 'bgc'
        ]
        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
236
237
        assert lang in model_urls[
            'rec'], 'param lang must in {}, but got {}'.format(
WenmuZhou's avatar
WenmuZhou committed
238
                model_urls['rec'].keys(), lang)
tink2123's avatar
tink2123 committed
239
240
241
242
        if lang == "ch":
            det_lang = "ch"
        else:
            det_lang = "en"
WenmuZhou's avatar
WenmuZhou committed
243
        use_inner_dict = False
WenmuZhou's avatar
WenmuZhou committed
244
        if postprocess_params.rec_char_dict_path is None:
WenmuZhou's avatar
WenmuZhou committed
245
            use_inner_dict = True
WenmuZhou's avatar
WenmuZhou committed
246
247
            postprocess_params.rec_char_dict_path = model_urls['rec'][lang][
                'dict_path']
WenmuZhou's avatar
WenmuZhou committed
248

249
250
        # init model dir
        if postprocess_params.det_model_dir is None:
tink2123's avatar
tink2123 committed
251
252
            postprocess_params.det_model_dir = os.path.join(BASE_DIR, VERSION,
                                                            'det', det_lang)
253
        if postprocess_params.rec_model_dir is None:
tink2123's avatar
tink2123 committed
254
255
            postprocess_params.rec_model_dir = os.path.join(BASE_DIR, VERSION,
                                                            'rec', lang)
WenmuZhou's avatar
WenmuZhou committed
256
        if postprocess_params.cls_model_dir is None:
tink2123's avatar
tink2123 committed
257
            postprocess_params.cls_model_dir = os.path.join(BASE_DIR, 'cls')
258
        print(postprocess_params)
WenmuZhou's avatar
WenmuZhou committed
259
        # download model
tink2123's avatar
tink2123 committed
260
261
        maybe_download(postprocess_params.det_model_dir,
                       model_urls['det'][det_lang])
WenmuZhou's avatar
WenmuZhou committed
262
263
264
        maybe_download(postprocess_params.rec_model_dir,
                       model_urls['rec'][lang]['url'])
        maybe_download(postprocess_params.cls_model_dir, model_urls['cls'])
WenmuZhou's avatar
WenmuZhou committed
265
266
267
268
269
270
271

        if postprocess_params.det_algorithm not in SUPPORT_DET_MODEL:
            logger.error('det_algorithm must in {}'.format(SUPPORT_DET_MODEL))
            sys.exit(0)
        if postprocess_params.rec_algorithm not in SUPPORT_REC_MODEL:
            logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
            sys.exit(0)
WenmuZhou's avatar
WenmuZhou committed
272
273
274
        if use_inner_dict:
            postprocess_params.rec_char_dict_path = str(
                Path(__file__).parent / postprocess_params.rec_char_dict_path)
WenmuZhou's avatar
WenmuZhou committed
275
276
277
278

        # init det_model and rec_model
        super().__init__(postprocess_params)

279
    def ocr(self, img, det=True, rec=True, cls=True):
WenmuZhou's avatar
WenmuZhou committed
280
281
282
283
284
285
286
287
        """
        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
288
289
290
        if isinstance(img, list) and det == True:
            logger.error('When input a list of images, det must be false')
            exit(0)
291
        if cls == True and self.use_angle_cls == False:
WenmuZhou's avatar
WenmuZhou committed
292
293
294
            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
295

WenmuZhou's avatar
WenmuZhou committed
296
        if isinstance(img, str):
WenmuZhou's avatar
WenmuZhou committed
297
298
299
300
            # download net image
            if img.startswith('http'):
                download_with_progressbar(img, 'tmp.jpg')
                img = 'tmp.jpg'
WenmuZhou's avatar
WenmuZhou committed
301
302
303
            image_file = img
            img, flag = check_and_read_gif(image_file)
            if not flag:
304
305
306
                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
307
308
309
            if img is None:
                logger.error("error in loading image:{}".format(image_file))
                return None
WenmuZhou's avatar
WenmuZhou committed
310
311
        if isinstance(img, np.ndarray) and len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
WenmuZhou's avatar
WenmuZhou committed
312
        if det and rec:
313
            dt_boxes, rec_res = self.__call__(img, cls)
WenmuZhou's avatar
WenmuZhou committed
314
315
316
317
318
319
320
321
322
            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]
323
            if self.use_angle_cls and cls:
WenmuZhou's avatar
WenmuZhou committed
324
325
326
                img, cls_res, elapse = self.text_classifier(img)
                if not rec:
                    return cls_res
WenmuZhou's avatar
WenmuZhou committed
327
328
            rec_res, elapse = self.text_recognizer(img)
            return rec_res
329
330
331


def main():
WenmuZhou's avatar
WenmuZhou committed
332
    # for cmd
333
    args = parse_args_whl(mMain=True)
WenmuZhou's avatar
WenmuZhou committed
334
335
336
337
338
339
    image_dir = args.image_dir
    if image_dir.startswith('http'):
        download_with_progressbar(image_dir, 'tmp.jpg')
        image_file_list = ['tmp.jpg']
    else:
        image_file_list = get_image_file_list(args.image_dir)
340
341
342
    if len(image_file_list) == 0:
        logger.error('no images find in {}'.format(args.image_dir))
        return
WenmuZhou's avatar
WenmuZhou committed
343
344

    ocr_engine = PaddleOCR(**(args.__dict__))
345
    for img_path in image_file_list:
WenmuZhou's avatar
WenmuZhou committed
346
347
348
349
350
351
352
353
        logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
        result = ocr_engine.ocr(img_path,
                                det=args.det,
                                rec=args.rec,
                                cls=args.use_angle_cls)
        if result is not None:
            for line in result:
                logger.info(line)