Ver1.0提交

.idea/DemoProject02.iml

@@ -2,7 +2,7 @@
 <module type="PYTHON_MODULE" version="4">
   <component name="NewModuleRootManager">
     <content url="file://$MODULE_DIR$" />
-    <orderEntry type="inheritedJdk" />
+    <orderEntry type="jdk" jdkName="Python 3.11" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

.idea/misc.xml

@@ -3,5 +3,5 @@
   <component name="Black">
     <option name="sdkName" value="Python 3.10 (PyCharmEnv)" />
   </component>
-  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.10 (PyCharmEnv)" project-jdk-type="Python SDK" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.11" project-jdk-type="Python SDK" />
 </project>

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

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

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('输入有误')

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