SNIC（Python实现）

# @Time : 2022/1/29
# @Author : hcl
import cv2
import numpy as np
import queue
import time
from matplotlib import pyplot as plt


def imgShow(window_title, image):
    cv2.imshow(window_title, image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def distance(center, element):
    """
    计算element和聚类中心的距离
    :param center: 聚类中心
    :param element:
    :return:
    """
    distance_x = pow(center[0] - element[0], 2) + pow(center[1] - element[1], 2)
    distance_c = pow(img_lab[center[0]][center[1]] - img_lab[element[0]][element[1]], 2)
    distance_c = sum(distance_c)
    return pow(distance_x / height_s + distance_c / m, 0.5)


def updateCenter():
    """
    更新质心
    :return:
    """
    # points是所有以key为质心的点的列表
    old_center = R[old_index]
    sum_points[old_index][0] += key[0]
    sum_points[old_index][1] += key[1]
    sum_points[old_index][2] += 1
    points = result[old_center]
    value = sum_points[old_index]
    new_center = value[0] // value[2], value[1] // value[2]
    R[old_index] = new_center
    result.pop(old_center)
    points.append(key)
    result[new_center] = points


def findNeighbor(point):
    neighbors = [(point[0] - 1, point[1] - 1), (point[0] - 1, point[1]), (point[0] - 1, point[1] + 1),
                 (point[0], point[1] - 1), (point[0], point[1] + 1),
                 (point[0] + 1, point[1] - 1), (point[0] + 1, point[1]), (point[0] + 1, point[1] + 1)]
    new_neighbors = []
    for x, y in neighbors:
        if 0 <= x < height and 0 <= y < width:
            new_neighbors.append((x, y))
    return new_neighbors


start_time = time.time()
img = cv2.imread(r'161062.jpg')
# 如果图片过大可以先resize
# img = cv2.resize(img, (0, 0), fx=0.2, fy=0.2)
print(img.shape)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
img_lab = np.array(img_lab, dtype=int)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
height, width, _ = img_lab.shape
num_points = height * width
# height_s和width_s是初始化一个簇的大小，我在这里初始化的时候没有初始化簇心的数量，我初始化的是簇的大小
height_s = height // 10
width_s = width // (width // height_s)
print(width_s, height_s)
num_k = height // height_s * (width // width_s)
print(f'簇心个数：{num_k}')
m = 20
# 点 和 index的关系
L = {}
# index 和 点的关系
R = {}
# 点: [点的横坐标之和, 点的纵坐标之和, 点的数量]
sum_points = {}
elements = queue.PriorityQueue()
result = {}
for i in range(height):
    for j in range(width):
        L[(i, j)] = 0

index = 0
min_dist = {}
for i in range(height // height_s):
    for j in range(width // width_s):
        index += 1
        R[index] = (i * height_s + height_s // 2, j * width_s + width_s // 2)
        elements.put([0, R[index], index])
        min_dist[R[index]] = 0
        result[(i * height_s + height_s // 2, j * width_s + width_s // 2)] = []
        sum_points[index] = [0, 0, 0]

i = 1
rate = 0
while not elements.empty():
    e = elements.get()
    key = e[1]
    if L[key] == 0:
        if i / num_points * 100 > rate + 0.1:
            print('\r', end='')
            rate = i / num_points * 100
            print(f'进度：{rate}%', end='')
        i += 1
        L[key] = e[2]
        old_index = L[key]
        updateCenter()
        key_neighbors = findNeighbor(key)
        for key_neighbor in key_neighbors:
            if L[key_neighbor] == 0:
                dist = distance(R[old_index], key_neighbor)
                if key_neighbor not in min_dist:
                    min_dist[key_neighbor] = dist
                    elements.put([dist, key_neighbor, old_index])
                elif dist < min_dist[key_neighbor]:
                    min_dist[key_neighbor] = dist
                    elements.put([dist, key_neighbor, old_index])

j = 0
for k, v in result.items():
    for i in v:
        gray[i[0]][i[1]] = 31 * j
    j += 1

# 绘制边界
for i in range(1, height - 1):
    for j in range(1, width - 1):
        count = 0
        point_neighbor = [(i - 1, j - 1), (i - 1, j), (i - 1, j + 1),
                          (i, j - 1), (i, j + 1),
                          (i + 1, j - 1), (i + 1, j), (i + 1, j + 1)]
        for neighbor in point_neighbor:
            if gray[neighbor[0]][neighbor[1]] != gray[i][j]:
                count += 1
        if count > 2:
            img[i][j][0] = 255
            img[i][j][1] = 0
            img[i][j][2] = 0

print('\r', end='')
print('进度：100%')
print(f'总耗时：{time.time() - start_time}秒')
plt.imshow(img)
plt.show()