import cv2
from paddleocr import PaddleOCR
from matplotlib import pyplot as plt
import os

#https://blog.csdn.net/qq_54827663/article/details/128067229?ops_request_misc=%257B%2522request%255Fid%2522%253A%252244852033-C534-4B01-A2D8-A0B994E869EB%2522%252C%2522scm%2522%253A%252220140713.130102334..%2522%257D&request_id=44852033-C534-4B01-A2D8-A0B994E869EB&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~top_positive~default-2-128067229-null-null.142^v100^pc_search_result_base6&utm_term=python%E8%BD%A6%E7%89%8C%E8%AF%86%E5%88%AB%E4%BB%A3%E7%A0%81&spm=1018.2226.3001.4187
"""
1. 转灰度图
由于车牌颜色多种多样，无法使用颜色阈值的方法来提取车牌，因此考虑使用形态学操作的方法来提取车牌。
首先将图片转为灰度图，排除颜色对识别的干扰。
"""
img = cv2.imread("C:/phone/016.jpg")
img = cv2.resize(img, (int(img.shape[1] * 0.5), int(img.shape[0] * 0.5)))
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

fig = plt.figure(figsize=(6, 6))
plt.imshow(gray, "gray"), plt.axis('off'), plt.title("gray")
plt.show()

"""
2. 顶帽运算
图像顶帽运算->原始图像减去图像开运算的结果，得到图像的噪声。
"""
# 创建一个17*17矩阵内核
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (17, 17))
tophat = cv2.morphologyEx(gray, cv2.MORPH_TOPHAT, kernel)

fig = plt.figure(figsize=(6, 6))
plt.imshow(tophat, "gray"), plt.axis('off'), plt.title("tophat")
plt.show()

"""
 3. Sobel算子提取y方向边缘
使用Sobel算子对字符进行y方向提取，经测试y方向的提取效果优于x方向。
"""
y = cv2.Sobel(tophat, cv2.CV_16S, 1, 0)
absY = cv2.convertScaleAbs(y)
fig = plt.figure(figsize=(6, 6))

plt.imshow(absY, "gray"), plt.axis('off'), plt.title("absY")
plt.show()

"""
4. 自适应二值化
灰度图像二值化，将灰度小于75的像素置为0，大于等于75的像素置为255。
"""
ret, binary = cv2.threshold(absY, 75, 255, cv2.THRESH_BINARY)

fig = plt.figure(figsize=(6, 6))
plt.imshow(binary, "gray"), plt.axis('off'), plt.title("binary")
plt.show()

"""
5. 开运算分割（纵向去噪，分隔）
使用形状为（1，15）的矩形kernel对图像进行Y方向开运算，使图像先腐蚀后膨胀，从而消除图像中的小面积白点。
"""
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 15))
Open = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)

fig = plt.figure(figsize=(6, 6))
plt.imshow(Open, "gray"), plt.axis('off'), plt.title("Open")
plt.show()


"""
6. 闭运算合并
使用形状为（41，15）的矩形kernel对图像进行偏 X方向的闭运算，将图像进行X方向融合找出车牌区域。
"""
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (41, 15))
close = cv2.morphologyEx(Open, cv2.MORPH_CLOSE, kernel)

fig = plt.figure(figsize=(6, 6))
plt.imshow(close, "gray"), plt.axis('off'), plt.title("close")
plt.show()

"""
7. 膨胀/腐蚀
再次使用特定大小的kernel对图像进行膨胀，腐蚀操作去噪，得到车牌区域。
"""
# 中远距离车牌识别
kernel_x = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 7))
kernel_y = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 11))

# 7-1、腐蚀、膨胀（去噪）
erode_y = cv2.morphologyEx(close, cv2.MORPH_ERODE, kernel_y)
dilate_y = cv2.morphologyEx(erode_y, cv2.MORPH_DILATE, kernel_y)

# 7-2、膨胀、腐蚀（连接）（二次缝合）
dilate_x = cv2.morphologyEx(dilate_y, cv2.MORPH_DILATE, kernel_x)
erode_x = cv2.morphologyEx(dilate_x, cv2.MORPH_ERODE, kernel_x)

fig = plt.figure(figsize=(15, 15))
plt.subplot(131), plt.imshow(erode_y, "gray"), plt.axis('off'), plt.title("erode_y")
plt.subplot(132), plt.imshow(dilate_y, 'gray'), plt.axis('off'), plt.title("dilate_y")
plt.subplot(133), plt.imshow(erode_x, 'gray'), plt.axis('off'), plt.title("erode_x")
plt.show()

"""
 8. 腐蚀、膨胀：去噪
再次使用特定大小的kernel对图像进行膨胀，腐蚀操作去除噪声保留车牌区域。
"""
kernel_e = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 9))
erode = cv2.morphologyEx(erode_x, cv2.MORPH_ERODE, kernel_e)

kernel_d = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 11))
dilate = cv2.morphologyEx(erode, cv2.MORPH_DILATE, kernel_d)

fig = plt.figure(figsize=(10, 10))
plt.subplot(121), plt.imshow(erode, "gray"), plt.axis('off'), plt.title("erode")
plt.subplot(122), plt.imshow(dilate, 'gray'), plt.axis('off'), plt.title("dilate")
plt.show()

"""
9. 获取外轮廓
对二值图进行轮廓提取。
"""
img_copy = img.copy()

# 9-1、得到轮廓
contours, hierarchy = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 9-2、画出轮廓并显示
cv2.drawContours(img_copy, contours, -1, (255, 0, 255), 2)

# cv2.imshow('Contours', img_copy)
fig = plt.figure(figsize=(6, 6))
plt.imshow(img_copy, "gray"), plt.axis('off'), plt.title("Contours")
plt.show()

"""
10. 遍历所有轮廓，找到车牌轮廓
使用矩形拟合所有轮廓，将满足：高*3<宽<高*7的矩形轮廓表示的ROI作为车牌区域返回。
"""
count = 0
ROI = None
for contour in contours:
    area = cv2.contourArea(contour)  # 计算轮廓内区域的面积
    # 得到矩形区域：左顶点坐标、宽和高
    x, y, w, h = cv2.boundingRect(contour)  # 获取坐标值和宽度、高度

    # 判断宽高比例、面积，截取符合图片
    if h * 3 < w < h * 7 and area > 1000:
        # 截取车牌并显示
        ROI = img[(y - 5):(y + h + 5), (x - 5):(x + w + 5)]  # 高，宽
        try:
            count += 1
            fig = plt.figure(figsize=(6, 6))
            plt.imshow(ROI), plt.axis('off'), plt.title("ROI")
            plt.show()
            # return ROI

        except:
            print("ROI提取出错！")
            # return
            pass

# 最后一步识别区
    os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # 设置允许重复加载动态链接库，若不允许，使用jupyter运行时内核会挂掉
    # Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
    # 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
    ocr = PaddleOCR(use_angle_cls=False, use_gpu=False,
                    lang="ch", show_log=False)  # need to run only once to download and load model into memory
    if ROI is None:
        print("没有提取到车牌")
        exit()

    fig = plt.figure(figsize=(6, 6))
    plt.imshow(ROI), plt.axis('off'), plt.title("ROI")
    plt.show()

    ocr_text = ocr.ocr(ROI, cls=False)  # 进行ROI识别
    for line in ocr_text:
        number_plate = line[-1][-1][0]
        print("车牌：", end="")
        print(number_plate)