灰度图像的K-means的python代码实现

import numpy as np
import cv2
import matplotlib.pyplot as plt
def kmeans(image, num_clusters, seed=0, max_iter=10000):
    np.random.seed(seed)
    img = image.reshape(-1, 1)
    #print(max(img))
    index = 0
    
    clus_value = np.array(np.random.rand(num_clusters)*max(img), dtype=float)
    #print(clus_value[0])

    while True:
        index += 1
        flag = clus_value.copy()

        cs = np.array([np.square(img-clus_value[0]), np.square(img-clus_value[1])])

        labels = np.argmin(cs, axis=0)

        for j in range(num_clusters):


            clus_value[j] = np.mean(img[labels==j])
        if np.sum(np.abs(clus_value-flag)) < 1e-8 or index==max_iter:
            break
    
    segmented_image = np.array(clus_value[labels].reshape(image.shape), dtype=np.uint16)
    return segmented_image
img = cv2.imread("image.tiff",2)
segment_image1 = np.zeros(img.shape,dtype="uint16")
segment_image = kmeans(img, num_clusters=2)
print(segment_image.shape)
for i in range(img.shape[0]):
    for j in range(img.shape[1]):
        if segment_image[i][j] == np.max(segment_image):
            segment_image1[i][j] = 65535
cv2.imwrite("result.tiff",segment_image1)
plt.imshow(segment_image1)
plt.show()