- import cv2
- import numpy as np
- from sklearn.cluster import KMeans
- import time
-
- # 中文路径读取
- def cv_imread(filePath, cv2_falg=cv2.COLOR_BGR2RGB):
- cv_img = cv2.imdecode(np.fromfile(filePath, dtype=np.uint8), cv2_falg)
- return cv_img
-
- # 自定义装饰器计算时间
- def compute_time(func):
- def compute(*args, **kwargs):
- st = time.time()
- result = func(*args, **kwargs)
- et = time.time()
- print('消费时间 %.6f s' % (et - st))
- return result
-
- return compute
-
-
-
- @compute_time
- def kmeans_img(image, num_clusters, show=False):
- # 如果图像是灰度图(单通道),将其转换为三通道
- if len(image.shape) == 2:
- image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-
- # 将图像的形状进行调整以便进行 K 均值聚类,提高训练速度
- pixels = cv2.resize(image.copy(), None, fx=0.05, fy=0.05, interpolation=cv2.INTER_LINEAR)
- pixels = np.float32(pixels.reshape((-1, 3)))
-
- segmented_pixels = np.float32(image.reshape((-1, 3)))
-
- # 初始化 KMeans 模型并拟合数据
- kmeans = KMeans(n_clusters=num_clusters)
- kmeans.fit(pixels)
-
- # 获取每个像素所属的簇标签
- labels = kmeans.predict(segmented_pixels)
-
- # 根据簇标签,将图像像素值转换为簇中心值
- segmented_image = kmeans.cluster_centers_[labels]
- segmented_image = np.uint8(segmented_image.reshape(image.shape))
-
- if show:
- plt.figure(figsize=(10, 5))
-
- plt.subplot(1, 2, 1)
- plt.title('Original Image')
- plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
- plt.axis('off')
-
- plt.subplot(1, 2, 2)
- plt.title('Segmented Image')
- plt.imshow(segmented_image)
- plt.axis('off')
-
- plt.tight_layout()
- plt.show()
-
- return segmented_image
- image_path =r"C:\Users\pc\Pictures\test\快.png"
- image = cv_imread(image_path)
- kmeans_img(image,4, show=True)

使用opencv内设的kmeans函数:直接原图进行训练,然后获取每个像素点的类,速度慢。上述方法对图像进行一个缩放后,训练模型,然后用模型再预测原图的每个像素点,速度快。
- def kmeans_img(image, num_clusters, show=True):
- # 如果图像是灰度图(单通道),将其转换为三通道
- if len(image.shape) == 2:
- image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
- # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
- print(image.shape)
- # 将图像的形状进行调整以便进行 K 均值聚类
- pixels = image.reshape((-1, 3))
- pixels = np.float32(pixels)
-
-
- # 设定 kmeans 参数并运行算法
- criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
-
- _, labels, centers = cv2.kmeans(pixels, num_clusters, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
-
- # 将图像像素值转换为簇中心值
- centers = np.uint8(centers)
- segmented_image = centers[labels.flatten()]
- segmented_image = segmented_image.reshape(image.shape)
-
- if show:
- # 显示原始图像和分割后的图像
- plt.figure(figsize=(10, 5))
-
- plt.subplot(1, 2, 1)
- plt.title('Original Image')
- plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
- plt.axis('off')
-
- plt.subplot(1, 2, 2)
- plt.title('Segmented Image')
- plt.imshow(segmented_image)
- plt.axis('off')
-
- plt.tight_layout()
- plt.show()
- return segmented_image