二分KMeans实现图片分割代码
时间: 2024-02-05 17:13:24 浏览: 24
以下是二分KMeans实现图片分割的Python代码:
```python
import numpy as np
from PIL import Image
def load_image(path):
img = Image.open(path)
return np.array(img)
def save_image(array, path):
img = Image.fromarray(array)
img.save(path)
def distance(x, y):
return np.sqrt(np.sum((x-y)**2))
def kmeans(points, k, max_iter=100):
centers = points[np.random.choice(range(len(points)), k, replace=False)]
for _ in range(max_iter):
clusters = [[] for _ in range(k)]
for point in points:
distances = [distance(point, center) for center in centers]
cluster = np.argmin(distances)
clusters[cluster].append(point)
new_centers = []
for i in range(k):
new_center = np.mean(clusters[i], axis=0)
new_centers.append(new_center)
if np.allclose(centers, new_centers):
break
centers = new_centers
return centers, clusters
def binary_kmeans(image_path, k=2):
image = load_image(image_path)
points = np.reshape(image, (-1, 3))
centers, clusters = kmeans(points, k)
labels = np.zeros(len(points), dtype=int)
for i, cluster in enumerate(clusters):
for point in cluster:
labels[np.where((points == point).all(axis=1))] = i
segmented_image = np.reshape(labels, (image.shape[0], image.shape[1]))
return segmented_image
if __name__ == '__main__':
segmented_image = binary_kmeans('image.jpg', k=2)
save_image(segmented_image, 'segmented_image.png')
```
其中,`load_image`函数将图片读入为numpy数组,`save_image`函数将numpy数组保存为图片。`distance`函数计算两个点之间的欧氏距离。`kmeans`函数执行kmeans算法,返回聚类中心和每个点所属的聚类。`binary_kmeans`函数对图片进行二分kmeans分割,返回分割后的二值图像。