tuoheng_algN/util/OcrBaiduSdk.py

import time
import traceback
from os.path import join

import cv2
from aip import AipOcr
from loguru import logger

from common.YmlConstant import baidu_yml_path
from enums.BaiduSdkEnum import BAIDUERRORDATA
from enums.ExceptionEnum import ExceptionType
from exception.CustomerException import ServiceException
from util.RWUtils import getConfigs


class OcrBaiduSdk:

    __slots__ = ('client', "__config")

    def __init__(self, base_dir, env):
        self.client = None
        self.__config = getConfigs(join(base_dir, baidu_yml_path % env))
        self.init_client()

    def init_client(self):
        if self.client is None:
            self.client = AipOcr(str(self.__config["orc"]["APP_ID"]),
                                 self.__config["orc"]["API_KEY"],
                                 self.__config["orc"]["SECRET_KEY"])

    '''
    {
    "log_id": 2471272194,
    "words_result_num": 2,
    "words_result":
        [
            {"words": " TSINGTAO"},
            {"words": "青島睥酒"}
        ]
    }
    '''

    def universal_text_recognition(self, image, request_id):
        reply_num = 0
        reply_value = None
        while True:
            try:
                or_result, or_image = cv2.imencode(".jpg", image)
                options = {
                    "language_type": "CHN_ENG",
                    "detect_direction": "true",
                    "detect_language": "true",
                    "probability": "true"
                }
                res_image = self.client.basicGeneral(or_image.tobytes(), options)
                error_code = res_image.get("error_code")
                if error_code:
                    enum = BAIDUERRORDATA.get(error_code)
                    # 如果异常编码未知, 返回空值
                    if enum is None:
                        logger.error("百度云人流量统计异常！error_code:{}, request_id: {}", error_code, request_id)
                        return None
                    # 重试指定次数后，还是异常，输出统一内部异常
                    if enum.value[3] == 0:
                        if reply_value is None:
                            reply_value = enum.value[4]
                        logger.error("百度云人流量统计异常！error_code:{}, error_msg:{}, reply_num:{}, request_id: {}",
                                     enum.value[0], enum.value[2], reply_num, request_id)
                        raise Exception()
                    # 重试指定次数后，还是异常，输出对应的异常
                    if enum.value[3] == 1:
                        if reply_value is None:
                            reply_value = enum.value[4]
                        raise ServiceException(str(enum.value[0]), enum.value[2])
                    # 重试指定次数后，还是异常，输出空
                    if enum.value[3] == 2:
                        if reply_value is None:
                            reply_value = enum.value[4]
                        if reply_num >= reply_value:
                            return None
                        raise Exception()
                return res_image
            except Exception as e:
                time.sleep(1)
                reply_num += 0.5
                self.init_client()
                if reply_num > reply_value:
                    if isinstance(e, ServiceException):
                        raise ServiceException(e.code, e.msg)
                    logger.error("通用文字识别失败: {}, request_id: {}", traceback.format_exc(), request_id)
                    raise ServiceException(ExceptionType.SERVICE_INNER_EXCEPTION.value[0],
                                           ExceptionType.SERVICE_INNER_EXCEPTION.value[1])

    '''
    {
        "log_id": 3583925545,
        "words_result": {
            "color": "blue",
            "number": "苏HS7766"
        }
    }
    '''
    def license_plate_recognition(self, image, request_id):
        reply_num = 0
        reply_value = None
        while True:
            try:
                or_result, or_image = cv2.imencode(".jpg", image)
                res_image = self.client.licensePlate(or_image.tobytes(), {"multi_detect": "true"})
                error_code = res_image.get("error_code")
                if error_code:
                    enum = BAIDUERRORDATA.get(error_code)
                    # 如果异常编码未知, 返回空值
                    if enum is None:
                        logger.error("百度云人流量统计异常！error_code:{}, request_id: {}", error_code, request_id)
                        return None
                    # 重试指定次数后，还是异常，输出统一内部异常
                    if enum.value[3] == 0:
                        if reply_value is None:
                            reply_value = enum.value[4]
                        logger.error("百度云人流量统计异常！error_code:{}, error_msg:{}, reply_num:{}, request_id: {}",
                                     enum.value[0], enum.value[2], reply_num, request_id)
                        raise Exception()
                    # 重试指定次数后，还是异常，输出对应的异常
                    if enum.value[3] == 1:
                        if reply_value is None:
                            reply_value = enum.value[4]
                        raise ServiceException(str(enum.value[0]), enum.value[2])
                    # 重试指定次数后，还是异常，输出空
                    if enum.value[3] == 2:
                        if reply_value is None:
                            reply_value = enum.value[4]
                        if reply_num >= reply_value:
                            return None
                        raise Exception()
                return res_image
            except Exception as e:
                time.sleep(1)
                reply_num += 1
                self.init_client()
                if reply_num > reply_value:
                    if isinstance(e, ServiceException):
                        raise ServiceException(e.code, e.msg)
                    logger.error("车牌识别失败: {}, request_id: {}", traceback.format_exc(), request_id)
                    raise ServiceException(ExceptionType.SERVICE_INNER_EXCEPTION.value[0],
                                           ExceptionType.SERVICE_INNER_EXCEPTION.value[1])