import cv2
import numpy as np

'''
基于Opencv图像处理的车牌定位和分割
'''
def car_plate_recognize(car):
    """=========================== 1. 定位车牌（车牌检测）==========================="""
    # 蓝色、黄色和绿色所对应的色彩空间
    lower_blue = np.array([100, 110, 110])
    upper_blue = np.array([130, 255, 255])
    lower_yellow = np.array([15, 55, 55])
    upper_yellow = np.array([50, 255, 255])
    lower_green = np.array([35, 100, 100])
    upper_green = np.array([85, 255, 255])

    hsv = cv2.cvtColor(car, cv2.COLOR_BGR2HSV)  # 将BGR图像转化到HSV的颜色空间
    mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
    mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask_plate = cv2.bitwise_or(mask_blue, mask_yellow)
    mask_plate = cv2.bitwise_or(mask_plate, mask_green)

    # 根据阈值找到对应颜色
    mask = cv2.cvtColor(mask_plate, cv2.COLOR_GRAY2BGR)
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    Matrix = np.ones((20, 20), np.uint8)
    mask1 = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, Matrix)
    mask = cv2.morphologyEx(mask1, cv2.MORPH_OPEN, Matrix)  # 形态学开运算
    ret, mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)  # 二值化进而获取轮廓
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 获取轮廓 contours

    # 初始化 box
    box = None

    # 寻找轮廓最大的 定位车牌
    for i in range(len(contours)):
        cnt = contours[i]
        area = cv2.contourArea(cnt)
        if area > 3000:
            rect = cv2.minAreaRect(cnt)
            box = cv2.boxPoints(rect)
            box = np.int_(box)
            break  # 找到一个符合条件的就跳出循环

    if box is None:
        raise ValueError("No contours found that meet the size requirement.")

    plate = cv2.drawContours(car.copy(), [box], -1, (0, 255, 0), 3)

    """=========================== 2. 分割车牌中的每个字符 ==========================="""
    ys = [box[0, 1], box[1, 1], box[2, 1], box[3, 1]]
    xs = [box[0, 0], box[1, 0], box[2, 0], box[3, 0]]
    ys_sorted_index = np.argsort(ys)
    xs_sorted_index = np.argsort(xs)
    x1 = box[xs_sorted_index[0], 0]
    x2 = box[xs_sorted_index[3], 0]
    y1 = box[ys_sorted_index[0], 1]
    y2 = box[ys_sorted_index[3], 1]
    ROI_plate = plate[y1:y2, x1:x2]
    ROI_plate_gray = cv2.cvtColor(ROI_plate, cv2.COLOR_BGR2GRAY)  # 灰度化
    ROI_plate_blur = cv2.GaussianBlur(ROI_plate_gray, (5, 5), 0)  # 高斯滤波
    ret, ROI_plate_Binary = cv2.threshold(ROI_plate_blur, 127, 255, cv2.THRESH_BINARY)  # 二值化
    # 形态学腐蚀 去除边框
    kernel = np.ones((5, 5), dtype=np.uint8)
    ROI_erode = cv2.erode(ROI_plate_Binary, kernel, iterations=1)
    # 根据宽度 裁剪7个字符
    width = ROI_erode.shape[1]
    height = ROI_erode.shape[0]
    word_0 = ROI_erode[0:height, 0:np.uint8(height / 2)]
    word_1 = ROI_erode[0:height, np.uint8(height / 2):height]
    size = np.uint8((width - height) / 5)
    word_2 = ROI_erode[0:height, height + 0 * size:height + 1 * size]
    word_3 = ROI_erode[0:height, height + 1 * size:height + 2 * size]
    word_4 = ROI_erode[0:height, height + 2 * size:height + 3 * size]
    word_5 = ROI_erode[0:height, height + 3 * size:height + 4 * size]
    word_6 = ROI_erode[0:height, height + 4 * size:height + 5 * size]
    word_all = [word_0, word_1, word_2, word_3, word_4, word_5, word_6]
    return plate, word_all


if __name__ == "__main__":
    car = cv2.imread(r'/mnt/data/cropped_license_plate_0.jpg', 1)
    plate, _ = car_plate_recognize(car)
    cv2.imwrite("plate.jpg", plate)