Crazy Lazy Coder

import cv2 as cv
import numpy as np
import random
import math
import copy

roi = cv.imread('sunflower.png') # 이미지 불러오기

k = 6 # 분류될 군집의 갯수

Z = [] # 각 군집의 대표값
X = [] # 샘플 집합. 여기서는 이미지픽셀
Z_xlist = [] # 각 군집에 해당되는 샘플을 저장할 리스트
prev_list = [] # 리스트의 이전상태를 기록할 리스트

for i in range(roi.shape[0]):
    for j in range(roi.shape[1]):
        x = (roi.item(i,j,0),roi.item(i,j,1),roi.item(i,j,2),(i,j)) # 이미지의 각 픽셀의 rgb값+이미지 좌표로 저장
        X.append(x) # 샘플리스트에 저장

checked = [] # 좌표 중복을 방지하기 위한 리스트
for i in range(k): # K개의 군집 초기화
    while True: # 랜덤하게 한 좌표를 뽑아서 Z를 초기화
        iy = random.randint(0,roi.shape[0]-1)
        ix = random.randint(0,roi.shape[1]-1)
        if (iy,ix) not in checked: # 체크한 좌표가 아니면 기록후 무한반복 탈출
            checked.append((iy,ix))
            break
    pos = checked[-1] # 방금 찾은 좌표로 설정
    value = ((roi.item(pos[0],pos[1],0),roi.item(pos[0],pos[1],1),roi.item(pos[0],pos[1],2)),(0,0,0))
    # Z 값은 방금 찾은 좌표의 rgb값 + 각 rgb의 총합(초기에는 0)
    Z.append(value)
    Z_xlist.append([]) # 각 군집에 해당되는 x는 아직 없음

while True:
    for x in X: # 모든 샘플에 대해 검사
        z_list = [] # 각 군집 대표값과 샘플의 유사성을 저장할 리스트
        for z in Z: # 각 군집 대표값과 유사성 측정후 저장
            dist = math.sqrt((x[0] - z[0][0])**2 + (x[1] - z[0][1])**2 + (x[2] - z[0][2])**2)
            # 유클리드 거리로 색의 유사성을 측정했으므로 0에 가까울수록 유사함.
            z_list.append(dist)
        idx = z_list.index(min(z_list)) # 최소값이 가장 유사한 군집
        z_value, z_sum = Z[idx] # 군집 대표값, 총합 가져옴
        z_sum = (z_sum[0] + x[0],z_sum[1] + x[1],z_sum[2] + x[2]) # x의 픽셀값을 총합에 더함
        Z[idx] = (z_value,z_sum) # 군집 갱신. 총합만 달라짐
        Z_xlist[idx].append(x) # 샘플을 해당 군집에 저장
    if prev_list == Z_xlist: # 탈출조건; 군집리스트가 변화가 없으면 알고리즘을 중단.
        break
    prev_list = copy.deepcopy(Z_xlist) # 군집 리스트변화를 측정하기 위해 현재 리스트 기록
    for idx, z in enumerate(Z): # 각 Z값 갱신
        z_value, z_sum = z
        z_value = (z_sum[0] / len(Z_xlist[idx]),
        z_sum[1] / len(Z_xlist[idx]),
        z_sum[2] / len(Z_xlist[idx])
        ) # 대표값은 각 채널값의 평균.
        z = (z_value, (0,0,0)) # 다시 샘플들을 검사해야하므로 총합은 0이고 대표값만 수정
    Z_xlist.clear() # 샘플리스트 초기화
    for i in range(k):
        Z_xlist.append([])

result = np.zeros((roi.shape[0],roi.shape[1],roi.shape[2]),dtype=np.uint8) # 결과이미지

for idx, xlist in enumerate(Z_xlist):
    rand = random.Random()
    rand.seed(idx*10) # 일정한 색 출력을 위해 시드 고정. 시드값은 K의 인덱스 * 10
    b = rand.randint(0,255)
    g = rand.randint(0,255)
    r = rand.randint(0,255)
    for x in xlist: # 해당 군집의 샘플의 위치를 알아내서 군집색으로 표시
        i, j = x[3] 
        result.itemset(i,j,0,b)
        result.itemset(i,j,1,g)
        result.itemset(i,j,2,r)

cv.imshow('roi',roi) # 기존 이미지
cv.imshow('res',result) # 군집화된 이미지
cv.waitKey(0)
cv.destroyAllWindows()

K-Means 알고리즘을 적용하여 컬러영상을 분할하는 프로그램이다.

6개의 군집으로 이미지 화소들을 분할한다.