# -*- coding: utf-8 -*-
# @Time    : 2024/2/18 14:45
# @Author  : len
# @File    : find_qr_block_2.py
# @Software: PyCharm 
# @Comment :
import os

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


def qr_KMeans(image):

    # 计算中心区域的坐标
    height, width = image.shape[:2]
    center_region = image[height // 4: 3 * height // 4, width // 4: 3 * width // 4]

    # 将中心区域的颜色空间从BGR转换到HSV
    center_hsv = cv2.cvtColor(center_region, cv2.COLOR_BGR2HSV)

    # 将中心区域图像数据重塑为二维数组 pixels (center_height*center_width, num_channels)
    center_pixels = center_hsv.reshape((-1, 3))

    # 应用K-means聚类来找到中心区域的两种主要颜色
    kmeans = KMeans(n_clusters=2, random_state=0, n_init=10).fit(center_pixels)

    dominant_colors = kmeans.cluster_centers_.astype(int)
    # for color in dominant_colors:
    #     print(color)

    # 确定前景色的颜色范围
    # 这里我们选择距离黑色(或白色)最远的颜色作为前景色
    foreground_color = max(dominant_colors, key=lambda c: cv2.norm(c, cv2.NORM_L2))

    # 转换整个图像到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)    # 黑白取反

        # kernel = np.ones((5, 5), np.uint8)
        # erosion = cv2.erode(mask_inv, kernel, iterations=1)
        # dilation = cv2.dilate(erosion, kernel, iterations=1)
        # erosion = cv2.erode(dilation, kernel, iterations=1)

        # 显示掩码
        idx += 2
        axs[idx].imshow(mask_inv, cmap='gray')
        axs[idx].set_title('Mask of Color: {}'.format(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"./android_img/qr_kmeans_img"
# image_files = [f for f in os.listdir(input_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
#
# for image_path in tqdm(image_files):
#     image_full_path = os.path.join(input_path, image_path)
#     image = cv2.imread(image_full_path)
#     if image is not None:
#         qr_KMeans(image)
#     else:
#         print(f"Failed to load image: {image_path}")


input_path = r"./data/img_1.png"
# input_path = r"./data/img.png"
image = cv2.imread(input_path)
qr_KMeans(image)