84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
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)
|