import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans


def qr_KMeans(image):
    # 获取图像的高度和宽度
    height, width = image.shape[:2]

    # 将图像分为四部分，取右上角区域 (高度的一半和宽度的一半)
    right_top_region = image[0:height // 2, width // 2:]

    # 将右上角区域的颜色空间从BGR转换到HSV
    right_top_hsv = cv2.cvtColor(right_top_region, cv2.COLOR_BGR2HSV)

    # 将右上角区域图像数据重塑为二维数组 pixels (right_top_height * right_top_width, num_channels)
    right_top_pixels = right_top_hsv.reshape((-1, 3))

    # 使用KMeans聚类来找到右上角区域的两种主要颜色
    kmeans = KMeans(n_clusters=2, random_state=0, n_init=10).fit(right_top_pixels)
    dominant_colors = kmeans.cluster_centers_.astype(int)

    # 打印颜色以供调试
    for color in dominant_colors:
        print(f"Dominant Color: {color}")

    # 转换整个图像到HSV空间
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    fig, axs = plt.subplots(1, len(dominant_colors) + 2, figsize=(20, 6))

    # 在第一个位置显示原图
    axs[0].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    axs[0].set_title('Original Image with Contours')
    axs[0].axis('off')
    axs[0].set_xticks([])
    axs[0].set_yticks([])

    for idx, color in enumerate(dominant_colors):
        # 创建HSV颜色范围
        lower_val = np.clip(color - np.array([10, 30, 30]), 0, 255)
        upper_val = np.clip(color + np.array([10, 255, 255]), 0, 255)

        # 创建颜色范围内的掩码
        mask = cv2.inRange(hsv_image, lower_val, upper_val)

        # 寻找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 在原图上画出轮廓
        if idx == 0:
            cv2.drawContours(image, contours, -1, (255, 0, 0), 2)  # 蓝色
        else:
            cv2.drawContours(image, contours, -1, (0, 255, 0), 2)  # 绿色

        # 取反显示掩码
        mask_inv = cv2.bitwise_not(mask)

        # 显示掩码
        idx += 2
        axs[idx].imshow(mask_inv, cmap='gray')
        axs[idx].set_title(f'Mask of Color: {color}')
        axs[idx].axis('off')
        axs[idx].set_xticks([])
        axs[idx].set_yticks([])

    # 在第二个位置显示带有轮廓的原图
    axs[1].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    axs[1].set_title('Original Image with Contours')
    axs[1].axis('off')
    axs[1].set_xticks([])
    axs[1].set_yticks([])

    # 显示结果
    plt.tight_layout()
    plt.show()


# 载入图片
input_path = r"../data/img_1.png"
image = cv2.imread(input_path)
qr_KMeans(image)