import cv2
import numpy as np
from sklearn.cluster import KMeans
from PIL import Image, ImageOps
import matplotlib.pyplot as plt

# 定义函数，透视变换
import cv2
import numpy as np

import cv2
import numpy as np

def perspective_transform(img_path):
    # 通过cv2加载图片
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # 复制图片
    img_copy = img.copy()
    # 将图片转灰度图
    img_copy = cv2.cvtColor(img_copy, cv2.COLOR_BGR2GRAY)
    # 二值化
    # 使用OTSU阈值算法将灰度图转换为二值图
    _, binary_image = cv2.threshold(img_copy, 127, 255,  cv2.THRESH_OTSU)
    # inverted_binary_image = cv2.bitwise_not(binary_image)
    # 黑白颜色反转
    img_copy = cv2.bitwise_not(img_copy)

    plt.title("Original Image")
    plt.imshow(img_copy, cmap='gray')
    plt.show()

    wechatqr = cv2.wechat_qrcode.WeChatQRCode("wechat_qrcode/detect.prototxt", "wechat_qrcode/detect.caffemodel",
                                              "wechat_qrcode/sr.prototxt", "wechat_qrcode/sr.caffemodel")
    res, points = wechatqr.detectAndDecode(binary_image)
    # 在原图上标记二维码的角点
    if len(res) > 0:
        print(res)
        # for point in points[0]:
        #     cv2.circle(img, tuple(int(i) for i in point), 10, (255, 0, 0), -1)  # 红色圆点，半径为10

        point = points[0]
        width = max(np.linalg.norm(point[0] - point[1]), np.linalg.norm(point[2] - point[3]))
        height = max(np.linalg.norm(point[1] - point[2]), np.linalg.norm(point[3] - point[0]))
        # 不添加边距
        points_dst = np.array([[0, 0], [width, 0], [width, height], [0, height]], dtype=point.dtype)
        # 生成变换矩阵
        matrix = cv2.getPerspectiveTransform(point, points_dst)
        # 应用透视变换，纠正图像
        corrected_image = cv2.warpPerspective(img, matrix, (int(width), int(height)))

        return corrected_image, img  # 返回校正后的图像和标记后的原图

    return None, img  # 如果没有检测到二维码，仅返回原图

def perfect_reflectance_white_balance(img):
    # 将图像从 BGR 转换到 LAB 色彩空间
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)

    # 分离 L, A, B 通道
    l, a, b = cv2.split(lab)

    # 寻找图像中的最亮区域（假设这些是白色区域）
    max_lightness = np.max(l)
    L_white = np.where(l == max_lightness)

    # 计算这些区域在 A 和 B 通道上的平均值
    a_white_avg = np.mean(a[L_white])
    b_white_avg = np.mean(b[L_white])

    # 适当调整 A 和 B 通道的平均值，避免过度调整
    a_shift = (128 - a_white_avg) * 0.1  # 减少调整幅度
    b_shift = (128 - b_white_avg) * 0.1  # 减少调整幅度
    a = a.astype(np.float32) + a_shift
    b = b.astype(np.float32) + b_shift

    # 确保 a, b 范围在 0 到 255 之间，并转换回 uint8 类型
    a = np.clip(a, 0, 255).astype(np.uint8)
    b = np.clip(b, 0, 255).astype(np.uint8)

    # 重新合成 LAB 图像并转换回 BGR 色彩空间
    lab = cv2.merge([l, a, b])
    corrected_image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    return corrected_image


corrected_image, img = perspective_transform("IMG_3606.png")

# 判断corrected_image不为空
if corrected_image is not None:
    # 保存图片文件
    corrected_image_bgr = cv2.cvtColor(corrected_image, cv2.COLOR_RGB2BGR)
    cv2.imwrite("corrected_image.png", corrected_image_bgr)
    # 通过 PIL 显示图片
    plt.title("Original Image")
    plt.imshow(corrected_image, cmap='gray')
    plt.show()
else:
    print("目标无法识别")

# # 读取图像
# image = cv2.imread('corrected_image.png')
# # 应用完美反射白平衡算法
# white_balanced_image = perfect_reflectance_white_balance(image)
#
# # 显示结果
# cv2.imshow('Original Image', image)
# cv2.imshow('White Balanced Image', white_balanced_image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()