# -*- coding: utf-8 -*-
# @Time    : 2024/2/18 11:06
# @Author  : len
# @File    : 矫正车牌.py
# @Software: PyCharm 
# @Comment :

import cv2
import numpy as np

# 读取图像
image_path = './img/img_2.png'
image = cv2.imread(image_path)

# 检查图像是否正确加载
if image is None:
    raise ValueError("未能加载图像文件。")

# 转换为灰度图
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# 应用自适应阈值来突出车牌
thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)

# 查找轮廓
contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

# 车牌候选的轮廓
candidates = []
for cnt in contours:
    rect = cv2.minAreaRect(cnt)
    box = cv2.boxPoints(rect)
    box = np.intp(box)
    (x, y), (w, h), angle = rect
    if w > 0 and h > 0 and (w * h) > 1000:  # 面积要大于1000，过滤掉太小的轮廓
        aspect_ratio = max(w, h) / min(w, h)
        if 2 < aspect_ratio < 5:  # 车牌的宽高比大约在2到5之间
            candidates.append((box, angle))

# 对每个候选轮廓进行透视变换
for box, angle in candidates:
    # 调整角度以使车牌水平
    if angle < -45:
        angle += 90

    # 旋转原始图像，使车牌水平
    center = (int(image.shape[1] / 2), int(image.shape[0] / 2))
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, rotation_matrix, (image.shape[1], image.shape[0]))

    # 获取旋转后的灰度图像并重新寻找边缘
    rotated_gray = cv2.cvtColor(rotated, cv2.COLOR_BGR2GRAY)
    rotated_thresh = cv2.adaptiveThreshold(rotated_gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
    rotated_contours, _ = cv2.findContours(rotated_thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    # 重新定位车牌
    rotated_candidates = []
    for rotated_cnt in rotated_contours:
        rotated_rect = cv2.minAreaRect(rotated_cnt)
        rotated_box = cv2.boxPoints(rotated_rect)
        rotated_box = np.intp(rotated_box)
        (rx, ry), (rw, rh), rangle = rotated_rect
        if rw > 0 and rh > 0 and (rw * rh) > 1000:
            aspect_ratio_rotated = max(rw, rh) / min(rw, rh)
            if 2 < aspect_ratio_rotated < 5:
                rotated_candidates.append(rotated_box)

    # 对第一个候选框应用透视变换
    if rotated_candidates:
        box = rotated_candidates[0]
        width = int(rotated_rect[1][0])
        height = int(rotated_rect[1][1])

        # 确保宽度大于高度
        if width < height:
            width, height = height, width

        # 定义透视变换的目标坐标
        dst_pts = np.array([[0, height - 1],
                            [0, 0],
                            [width - 1, 0],
                            [width - 1, height - 1]], dtype="float32")
        src_pts = np.array(box, dtype="float32")

        # 计算透视变换矩阵
        M = cv2.getPerspectiveTransform(src_pts, dst_pts)

        # 应用透视变换
        warped = cv2.warpPerspective(rotated, M, (width, height))
        break

# 显示结果
cv2.imshow("Original Image", image)
cv2.imshow("Rotated and Corrected Plate", warped)
cv2.waitKey(0)
cv2.destroyAllWindows()