K-means法でいい感じにイラストを加工する

※これはOpenCVによるコンピュータビジョン・機械学習入門の演習問題そのまんまです。だから本見た方がはやいよ。

K-means法での減色処理がプログラム簡単な割に綺麗だったのでまとめようかなと思います。

K-means法

K-means法はクラスタリングの一種です。以下の手順でクラスタリングを行います。

• 初期化：データをランダムにK個のクラスに振り分け、それぞれのクラスの平均ベクトルを求める
  • 平均ベクトルを固定して各データがどの平均ベクトルに近いかを調べ、クラスラベルを更新する
  • クラスラベルを元に平均ラベルを更新する
  • 平均ベクトルとクラスラベルの変化が十分小さくなるまでこれを繰り返す

クラス数が決め打ちこととかは置いておいて理解しやすいから好きです。

pythonとOpenCVを使った減色処理の実装

今回はライブラリをばしばし使います。プログラムはこんな感じ。

import cv2
import numpy as np
import copy
import sys

file_src = sys.argv[1]

img_src = cv2.imread(file_src, 1)
img_shape = img_src.shape

pixels = np.reshape(img_src, (-1, 3))
pixels = np.float32(pixels)

criteria = (cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 100, 1.0)
flags = cv2.KMEANS_RANDOM_CENTERS
compactness, labels, centers = cv2.kmeans(pixels, 5, None, criteria, 10, flags)


result = copy.copy(img_src)
count = 0
for row in range(img_shape[0]):
    for col in range(img_shape[1]):
        result[row][col] = centers[labels[row*img_shape[1]+col]]

cv2.imwrite('result.png', result)

ちなみに手持ちのデータを入れてみた結果はこんな感じです。

クラス数5だとこんな感じになりました。結構きれい。

まとめ

K-means法で簡単にきれいな結果が出たのでまとめてみました。