"""
    模块作者：
        代码编写：焦雅雯
        代码优化/整合/打包：王昱博
    模块用途：
        使用OCR库进行车牌号识别和初步分类
"""

import cv2
import easyocr
import hyperlpr3 as lpr3


class OCR:
    @staticmethod
    def SwapChars(text: str) -> str:
        text = text.replace('I', '1')
        text = text.replace('O', '0')
        return text

    @classmethod
    def RecognizeLicensePlate(cls, image: cv2.Mat, lpr_text: str) -> tuple:
        reader = easyocr.Reader(['ch_sim', 'en'], model_storage_directory='./easyocr_model')
        result = reader.readtext(image)
        license_plate = ""
        for res in result:
            license_plate += res[-2]  # 如果车牌号码是两行的，按行识别出来再拼接起来
        license_plate = cls.SwapChars(license_plate)
        if lpr_text is not None:
            license_plate = lpr_text
        if '\u8b66' in license_plate:
            car_type = '警用车辆'
        elif '\u573a\u5185' in license_plate:
            car_type = '场内车辆'
        elif '\u6302' in license_plate:
            car_type = '挂车/半挂车'
        elif len(license_plate) > 7:
            car_type = '新能源车辆'
        else:
            car_type = '小型轿车'
        return license_plate, car_type

    @classmethod
    def RecognizeLicensePlate2(cls, image: cv2.Mat):
        reco = lpr3.LicensePlateCatcher()
        results = reco(image)
        for code, conf, _type, box in results:
            x0, y0, x1, y1 = box
            cut_image = image[y0:y1, x0:x1]
            return code, conf, cut_image
        return None, None, None