paddleocr.py

# 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
from ppocr.utils.utility import initial_logger

logger = initial_logger()
from ppocr.utils.utility import check_and_read_gif

__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'
    },
}

SUPPORT_DET_MODEL = ['DB']
SUPPORT_REC_MODEL = ['Rosetta', 'CRNN', 'STARNet', 'RARE']


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()
    if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
        logger.error("ERROR, something went wrong")
        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
    # 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


def parse_args():
    import argparse

    def str2bool(v):
        return v.lower() in ("true", "t", "1")

    parser = argparse.ArgumentParser()
    # params for prediction engine
    parser.add_argument("--use_gpu", type=str2bool, default=True)
    parser.add_argument("--ir_optim", type=str2bool, default=True)
    parser.add_argument("--use_tensorrt", type=str2bool, default=False)
    parser.add_argument("--gpu_mem", type=int, default=8000)

    # 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_max_side_len", type=float, default=960)

    # DB parmas
    parser.add_argument("--det_db_thresh", type=float, default=0.3)
    parser.add_argument("--det_db_box_thresh", type=float, default=0.5)
    parser.add_argument("--det_db_unclip_ratio", type=float, default=2.0)

    # EAST parmas
    parser.add_argument("--det_east_score_thresh", type=float, default=0.8)
    parser.add_argument("--det_east_cover_thresh", type=float, default=0.1)
    parser.add_argument("--det_east_nms_thresh", type=float, default=0.2)

    # 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_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(
        "--rec_char_dict_path",
        type=str,
        default="./ppocr/utils/ppocr_keys_v1.txt")
    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):
        """
        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)

        # 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

        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)

        postprocess_params.rec_char_dict_path = Path(
            __file__).parent / postprocess_params.rec_char_dict_path

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

    def ocr(self, img, det=True, rec=True):
        """
        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))
        if isinstance(img, str):
            image_file = img
            img, flag = check_and_read_gif(image_file)
            if not flag:
                img = cv2.imread(image_file)
            if img is None:
                logger.error("error in loading image:{}".format(image_file))
                return None
        if det and rec:
            dt_boxes, rec_res = self.__call__(img)
            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]
            rec_res, elapse = self.text_recognizer(img)
            return rec_res