diff --git a/.idea/DemoProject02.iml b/.idea/DemoProject02.iml index d0876a7..909438d 100644 --- a/.idea/DemoProject02.iml +++ b/.idea/DemoProject02.iml @@ -2,7 +2,7 @@ - + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml index 7d61941..aef59c4 100644 --- a/.idea/misc.xml +++ b/.idea/misc.xml @@ -3,5 +3,5 @@ - + \ No newline at end of file diff --git a/TestSet/中、大型汽车/川A802V6.jpg b/TestSet/中、大型汽车/川A802V6.jpg new file mode 100644 index 0000000..60daa88 Binary files /dev/null and b/TestSet/中、大型汽车/川A802V6.jpg differ diff --git a/TestSet/中、大型汽车/沪A33333.jpg b/TestSet/中、大型汽车/沪A33333.jpg new file mode 100644 index 0000000..cf84ae7 Binary files /dev/null and b/TestSet/中、大型汽车/沪A33333.jpg differ diff --git a/TestSet/军警、应急车辆/粤S3461.jpg b/TestSet/军警、应急车辆/粤S3461.jpg new file mode 100644 index 0000000..a83380c Binary files /dev/null and b/TestSet/军警、应急车辆/粤S3461.jpg differ diff --git a/TestSet/军警、应急车辆/苏C985A.jpg b/TestSet/军警、应急车辆/苏C985A.jpg new file mode 100644 index 0000000..85ecc94 Binary files /dev/null and b/TestSet/军警、应急车辆/苏C985A.jpg differ diff --git a/TestSet/场内车辆/粤AD0053.jpeg b/TestSet/场内车辆/粤AD0053.jpeg new file mode 100644 index 0000000..4b5869c Binary files /dev/null and b/TestSet/场内车辆/粤AD0053.jpeg differ diff --git a/TestSet/场内车辆/辽B04494.jpg b/TestSet/场内车辆/辽B04494.jpg new file mode 100644 index 0000000..8466f38 Binary files /dev/null and b/TestSet/场内车辆/辽B04494.jpg differ diff --git a/TestSet/外籍车/京A3429I.jpg b/TestSet/外籍车/京A3429I.jpg new file mode 100644 index 0000000..0a8df11 Binary files /dev/null and b/TestSet/外籍车/京A3429I.jpg differ diff --git a/TestSet/外籍车/鲁AL88888.jpg b/TestSet/外籍车/鲁AL88888.jpg new file mode 100644 index 0000000..17ce459 Binary files /dev/null and b/TestSet/外籍车/鲁AL88888.jpg differ diff --git a/TestSet/大型新能源车/粤AAKL8434.jpg b/TestSet/大型新能源车/粤AAKL8434.jpg new file mode 100644 index 0000000..df5a3c3 Binary files /dev/null and b/TestSet/大型新能源车/粤AAKL8434.jpg differ diff --git a/TestSet/大型新能源车/粤BBQ9397.jpg b/TestSet/大型新能源车/粤BBQ9397.jpg new file mode 100644 index 0000000..db21277 Binary files /dev/null and b/TestSet/大型新能源车/粤BBQ9397.jpg differ diff --git a/TestSet/小型新能源汽车/京AC11112.jpg b/TestSet/小型新能源汽车/京AC11112.jpg new file mode 100644 index 0000000..cd0c95b Binary files /dev/null and b/TestSet/小型新能源汽车/京AC11112.jpg differ diff --git a/TestSet/小型新能源汽车/沪AAR0112.jpg b/TestSet/小型新能源汽车/沪AAR0112.jpg new file mode 100644 index 0000000..1b0529f Binary files /dev/null and b/TestSet/小型新能源汽车/沪AAR0112.jpg differ diff --git a/TestSet/小型汽车/浙C99229.jpg b/TestSet/小型汽车/浙C99229.jpg new file mode 100644 index 0000000..5262716 Binary files /dev/null and b/TestSet/小型汽车/浙C99229.jpg differ diff --git a/TestSet/小型汽车/粵B99999.jpg b/TestSet/小型汽车/粵B99999.jpg new file mode 100644 index 0000000..386fb0c Binary files /dev/null and b/TestSet/小型汽车/粵B99999.jpg differ diff --git a/classification_ai.py b/classification_ai.py new file mode 100644 index 0000000..2e0f80f --- /dev/null +++ b/classification_ai.py @@ -0,0 +1,61 @@ +""" + 模块作者： + AI代码结构：刘钰廷、冯雅君 + 代码优化整理：王昱博、冯昌盛 + AI模型训练/纠错：刘钰廷、冯雅君、冯昌盛 + 代码整合/打包：王昱博 + 模块用途： + 图像分类AI，用于区分车牌的具体类型 +""" + +import cv2 +from PIL import Image +from pathlib import Path +from fastai.vision.all import * +from fastai.metrics import error_rate +from fastai.learner import load_learner +from torchvision.models import resnet34 +from fastai.vision.data import ImageBlock +from fastai.vision.core import imagenet_stats +from fastai.data.block import CategoryBlock, DataBlock +from fastai.vision.augment import Resize, aug_transforms +from fastai.vision.learner import cnn_learner, vision_learner +from fastai.data.transforms import get_image_files, parent_label, RandomSplitter, Normalize + + +class ClassificationAI: + @staticmethod + def ConvertImage(cv_img: cv2.Mat) -> Image.Image: + return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)).resize((460, 460)) + + @classmethod + def TrainAI(cls, data_set_path: str, export_path: str) -> None: + blocks = (ImageBlock, CategoryBlock) + batch_size = 32 + dls = DataBlock( + blocks=blocks, + get_items=get_image_files, + splitter=RandomSplitter(), + get_y=parent_label, + item_tfms=Resize(460), + batch_tfms=[*aug_transforms(size=224, min_scale=0.75), Normalize.from_stats(*imagenet_stats)] + ).dataloaders(data_set_path, num_workers=4, bs=batch_size) + model = vision_learner(dls, resnet34, metrics=error_rate) + model.fine_tune(5, freeze_epochs=3) # 5 - 训练的轮次， 3 - 冻结的轮次 + model.export(Path(export_path) / 'model.pkl') + + @classmethod + def PredictImage(cls, image: cv2.Mat, model_path: str) -> tuple: + # 加载模型 + model = load_learner(model_path) + # 读取图片并转换为Tensor + img = cls.ConvertImage(image) # 读取图像文件 + # 进行预测 + pred_class, pred_idx, outputs = model.predict(img) + # 获取置信度 + # 检查输出张量的维度 + if outputs.dim() == 0: + confidence = float(outputs) + else: + confidence = float(outputs[pred_idx]) + return pred_class, confidence diff --git a/classify_model/0.0625-2.pkl b/classify_model/0.0625-2.pkl new file mode 100644 index 0000000..b7fe389 Binary files /dev/null and b/classify_model/0.0625-2.pkl differ diff --git a/classify_model/0.0625.pkl b/classify_model/0.0625.pkl new file mode 100644 index 0000000..a7fe4f5 Binary files /dev/null and b/classify_model/0.0625.pkl differ diff --git a/classify_model/0.125.pkl b/classify_model/0.125.pkl new file mode 100644 index 0000000..b7b6edd Binary files /dev/null and b/classify_model/0.125.pkl differ diff --git a/cut_image.py b/cut_image.py new file mode 100644 index 0000000..e6cc372 --- /dev/null +++ b/cut_image.py @@ -0,0 +1,66 @@ +""" + 模块作者： + 图像预处理：潘浩宇 + 轮廓寻找与切分：戴晓齐 + 代码优化/整合/打包：王昱博 + 模块用途： + 对车牌图片进行预处理和切分，找出包含车牌号的部分 +""" + +import cv2 + +class ImageCutter: + @staticmethod + # 图像去噪灰度处理，消除噪点 + def gray_guss(image): + image = cv2.GaussianBlur(image, (3, 3), 0) + gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + return gray_image + + @classmethod + def ImagePreProcess(cls, image_path: str) -> tuple: + # 复制一张图片，在复制图上进行图像操作，保留原图 + origin_image = cv2.imread(image_path) + # 图像去噪灰度处理 + image = origin_image.copy() + # x方向上的边缘检测（增强边缘信息） + gray_image = cls.gray_guss(image) + Sobel_x = cv2.Sobel(gray_image, cv2.CV_16S, 1, 0) + absX = cv2.convertScaleAbs(Sobel_x) + image = absX + # 图像阈值化操作——获得二值化图，将像素置为0或者255。将灰度转成黑白 + ret, image = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU) + # 形态学（从图像中提取对表达和描绘区域形状有意义的图像分量）——闭操作 + # 使用形状为（30，10）的矩形kernelX对图像进行偏X方向的闭运算，将图像进行X方向融合找出车牌区域。 + kernelX = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 10)) + image = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernelX, iterations=1) + return origin_image, image + + @classmethod + def CutPlateRect(cls, origin_image: cv2.Mat, image: cv2.Mat) -> cv2.Mat: + # 去除细小的边缘 + # 腐蚀（erode）和膨胀（dilate） + kernelX = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 1)) + kernelY = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 20)) + # x方向进行闭操作（抑制暗细节） + image = cv2.dilate(image, kernelX) + image = cv2.erode(image, kernelX) + # y方向的开操作 + image = cv2.erode(image, kernelY) + image = cv2.dilate(image, kernelY) + # 中值滤波（去噪）将边缘平滑 + image = cv2.medianBlur(image, 21) + # 获得轮廓 RETR_EXTERNAL矩形的外边缘 + contours, hierarchy = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + # 筛选 + for item in contours: + rect = cv2.boundingRect(item) + x = rect[0] + y = rect[1] + weight = rect[2] + height = rect[3] + # 根据轮廓的形状特点，确定车牌的轮廓位置并截取图像 + if (weight > (height * 3.5)) and (weight < (height * 4)): # 对长宽比例进行确定 + _image = origin_image[y:y + height, x:x + weight] # 对图片进行裁剪 + return cv2.Mat(_image) + return origin_image diff --git a/easyocr_model/craft_mlt_25k.pth b/easyocr_model/craft_mlt_25k.pth new file mode 100644 index 0000000..7461d59 Binary files /dev/null and b/easyocr_model/craft_mlt_25k.pth differ diff --git a/easyocr_model/zh_sim_g2.pth b/easyocr_model/zh_sim_g2.pth new file mode 100644 index 0000000..e7d9dad Binary files /dev/null and b/easyocr_model/zh_sim_g2.pth differ diff --git a/main.py b/main.py index 8b13789..c257a41 100644 --- a/main.py +++ b/main.py @@ -1 +1,52 @@ +""" + 主程序作者：王昱博 + 车牌识别系统： + 使用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) + ocr_text, ocr_type = OCR.RecognizeLicensePlate(cut_image) + 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}') + cv2.waitKey(0) + + +if __name__ == '__main__': + result = input('请选择运行模式(训练(y)/识别(n)): ') + if result == 'y' or result == 'Y': + 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 == 'n' or result == 'N': + 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('输入有误') diff --git a/ocr.py b/ocr.py new file mode 100644 index 0000000..fdb9c15 --- /dev/null +++ b/ocr.py @@ -0,0 +1,46 @@ +""" + 模块作者： + 代码编写：焦雅雯 + 代码优化/整合/打包：王昱博 + 模块用途： + 使用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) -> 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] # 如果车牌号码是两行的，按行识别出来再拼接起来 + if '\u8b66' in license_plate: + car_type = 'police' + elif '\u573a\u5185' in license_plate: + car_type = 'internal' + elif '\u6302' in license_plate: + car_type = 'bigCar' + else: + car_type = 'smallCar' + 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, image diff --git a/车牌分类示例.zip b/车牌分类示例.zip new file mode 100644 index 0000000..6db560e Binary files /dev/null and b/车牌分类示例.zip differ