86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# @Time : 2024/10/16 上午9:54
|
||
|
|
# @Author : len
|
||
|
|
# @File : 003.py
|
||
|
|
# @Software: PyCharm
|
||
|
|
# @Comment :
|
||
|
|
|
||
|
|
|
||
|
|
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]
|
||
|
|
|
||
|
|
# 将整个图像的颜色空间从BGR转换到HSV
|
||
|
|
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||
|
|
|
||
|
|
# 将整个图像数据重塑为二维数组 pixels (height * width, num_channels)
|
||
|
|
pixels = hsv_image.reshape((-1, 3))
|
||
|
|
|
||
|
|
# 使用KMeans聚类来对整个图像进行二分类
|
||
|
|
kmeans = KMeans(n_clusters=2, random_state=0, n_init=10).fit(pixels)
|
||
|
|
dominant_colors = kmeans.cluster_centers_.astype(int)
|
||
|
|
|
||
|
|
# 打印颜色以供调试
|
||
|
|
for color in dominant_colors:
|
||
|
|
print(f"Dominant Color: {color}")
|
||
|
|
|
||
|
|
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)
|