|
|
"""
|
|
|
主程序作者：王昱博
|
|
|
|
|
|
车牌识别系统：
|
|
|
使用OCR技术对车牌号码进行识别
|
|
|
使用图像分类AI对车牌种类进行区分
|
|
|
"""
|
|
|
import cv2
|
|
|
|
|
|
from ocr import OCR
|
|
|
from cut_image import ImageCutter
|
|
|
from classification_ai import ClassificationAI
|
|
|
|
|
|
classify_models = ['.\\classify_model\\0.0625.pkl', '.\\classify_model\\0.0625-2.pkl', '.\\classify_model\\0.125.pkl']
|
|
|
|
|
|
|
|
|
def train(train_set_path: str, export_path: str) -> None:
|
|
|
ClassificationAI.TrainAI(train_set_path, export_path)
|
|
|
|
|
|
|
|
|
def main(classify_model_index: int, image_path: str) -> None:
|
|
|
global classify_models
|
|
|
origin_image, gray_image = ImageCutter.ImagePreProcess(image_path)
|
|
|
lpr_text, lpr_conf, cut_image = OCR.RecognizeLicensePlate2(origin_image)
|
|
|
if cut_image is None:
|
|
|
cut_image = ImageCutter.CutPlateRect(origin_image, gray_image)
|
|
|
ocr_text, ocr_type = OCR.RecognizeLicensePlate(cut_image, lpr_text)
|
|
|
if lpr_text is None:
|
|
|
lpr_text = ocr_text
|
|
|
lpr_conf = None
|
|
|
ai_type, ai_conf = ClassificationAI.PredictImage(cut_image, classify_models[classify_model_index])
|
|
|
print(f'识别完成，以下为识别结果：\n车牌号：{lpr_text} [置信度：{lpr_conf}]\n车牌类型：\n\t{ocr_type}(OCR推测)\n\t{ai_type}(AI分类识别)\n\tAI识别置信度：{ai_conf}')
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
result = input('请选择运行模式(训练(t)/识别(r)): ')
|
|
|
if result == 't' or result == 'T':
|
|
|
data_path = input('输入训练集路径: ')
|
|
|
export_path = input('输入模型保存路径: ')
|
|
|
try:
|
|
|
train(data_path, export_path)
|
|
|
except Exception as e:
|
|
|
print(f'训练过程中发生错误: {e}')
|
|
|
else:
|
|
|
print('模型已成功训练')
|
|
|
finally:
|
|
|
print('训练结束')
|
|
|
elif result == 'r' or result == 'R':
|
|
|
model_index = input('选择使用的识别模型(1/2/3): ')
|
|
|
image_path = input('输入图片路径: ')
|
|
|
if (not model_index.isdigit()) or (int(model_index) < 1) or (int(model_index) > 3):
|
|
|
print('输入有误')
|
|
|
else:
|
|
|
main(int(model_index), image_path)
|
|
|
else:
|
|
|
print('输入有误')
|