tuoheng_alg_move_20240715/util/OcrBaiduSdk.py

import time
from threading import Lock

import cv2
from aip import AipOcr
from loguru import logger

from enums.ExceptionEnum import ExceptionType
from exception.CustomerException import ServiceException


class OcrBaiduSdk:

    def __init__(self, content):
        self.content = content
        self.init_client()
        # self.lock = Lock()

    def init_client(self):
        self.client = AipOcr(str(self.content["baidu"]["orc"]["APP_ID"]), self.content["baidu"]["orc"]["API_KEY"],
                             self.content["baidu"]["orc"]["SECRET_KEY"])

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

    def universal_text_recognition(self, image, request_id):
        # try:
        #     self.lock.acquire()
        reply_num = 1
        while True:
            try:
                or_result, or_image = cv2.imencode(".jpg", image)
                res_image = self.client.basicGeneral(or_image.tobytes())
                if res_image.get("error_code") == 216630 or res_image.get("error_msg") == 'recognize error':
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    return None
                if res_image.get("error_code") == 282403 or res_image.get("error_msg") == 'target recognize error':
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    return None
                if res_image.get("error_code") == 216202 or res_image.get("error_msg") == 'image size error':
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    return None
                if res_image.get("error_code") is not None:
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    raise Exception("百度云调接口失败")
                return res_image
            except Exception as e:
                logger.exception("通用文字识别失败: {}, 当前重试次数：{}, request_id: {}", e, reply_num, request_id)
                time.sleep(1)
                reply_num += 1
                self.init_client()
                if reply_num > 5:
                    logger.exception("通用文字识别失败: {}, request_id: {}", e, request_id)
                    raise ServiceException(ExceptionType.UNIVERSAL_TEXT_RECOGNITION_FAILED.value[0],
                                           ExceptionType.UNIVERSAL_TEXT_RECOGNITION_FAILED.value[1])
        # except Exception as ee:
        #     logger.exception("通用文字识别加锁异常: {}, request_id: {}", ee, request_id)
        #     raise ServiceException(ExceptionType.UNIVERSAL_TEXT_RECOGNITION_FAILED.value[0],
        #                            ExceptionType.UNIVERSAL_TEXT_RECOGNITION_FAILED.value[1])
        # finally:
        #     self.lock.release()

    '''
    {
        "log_id": 3583925545,
        "words_result": {
            "color": "blue",
            "number": "苏HS7766"
        }
    }
    '''

    def license_plate_recognition(self, image, request_id):
        # try:
        #     self.lock.acquire()
        reply_num = 1
        while True:
            try:
                or_result, or_image = cv2.imencode(".jpg", image)
                res_image = self.client.licensePlate(or_image.tobytes())
                if res_image.get("error_code") == 282403 or res_image.get("error_msg") == 'target recognize error':
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    return None
                if res_image.get("error_code") == 216630 or res_image.get("error_msg") == 'recognize error':
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    return None
                if res_image.get("error_code") == 216202 or res_image.get("error_msg") == 'image size error':
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    return None
                if res_image.get("error_code") is not None:
                    logger.error("百度云调接口失败: {}, 当前重试次数：{}, request_id: {}", res_image, reply_num, request_id)
                    raise Exception("百度云调接口失败")
                return res_image
            except Exception as e:
                logger.exception("车牌识别失败: {}, 当前重试次数：{}, request_id: {}", e, reply_num, request_id)
                time.sleep(1)
                reply_num += 1
                self.init_client()
                if reply_num > 5:
                    logger.exception("车牌识别失败: {}, request_id: {}", e, request_id)
                    raise ServiceException(ExceptionType.ABNORMAL_LICENSE_PLATE_RECOGNITION.value[0],
                                           ExceptionType.ABNORMAL_LICENSE_PLATE_RECOGNITION.value[1])
        # except Exception as ee:
        #     logger.exception("车牌识别加锁异常: {}, request_id: {}", ee, request_id)
        #     raise ServiceException(ExceptionType.ABNORMAL_LICENSE_PLATE_RECOGNITION.value[0],
        #                            ExceptionType.ABNORMAL_LICENSE_PLATE_RECOGNITION.value[1])
        # finally:
        #     self.lock.release()