k-means法で画像を減色するサンプルコード

何かというと

下記のような画像ファイルの色数を減らして、N色の画像にする処理を考えます。どんな方法を思いつくでしょうか？（N=16あたりで考えてください。）

たとえば、次のような手順が考えられます。

　(1)画像の中で代表的な色をN個選ぶ。
　(2)それぞれのピクセルを一番色が近い代表色に置き換える。

(2)の処理については、「色の近さ」を定義する必要がありますが、これは (R,G,B)（R,G,Bはそれぞれ0〜255の値）で色を表現して、

$(R_1,G_1,B_1)$ と $(R_2,G_2,B_2)$ の距離
= $(R_1-R_2)^2+(G_1-G_2)^2+(B_1-B_2)^2$

という、3次元ベクトル空間上の距離で計算することができます。

一方、(1)において、どうやって代表色を選ぶかが悩ましいところです。機械学習のアルゴリズムの1つである「k-means法」を利用すると、このような代表色を自然に選択することができます。

これは次のようなアルゴリズムです。

(1) はじめに、ランダムにN個の代表色を選択します。
(2) 画像のそれぞれのピクセルについて、一番近い代表色を選んで、全ピクセルをN個のグループに分割します。
(3) それぞれのグループの平均色（R, G, Bそれぞれの平均値を取った色）を新たな代表色とします。
(4) 新しい代表色を用いて、再度、(2)の処理を行います。

このループを何度も繰り返すと、やがて、代表色が変化しなくなります（つまり、各グループの平均色が代表色と一致する）ので、これを最終的な代表色として採用します。

非常に単純なロジックなので、Pythonなどで簡単に実装することができます。

実装してみた

数値演算ライブラリのNumPyと画像処理ライブラリのpilを使っています。RHEL6/CentOS6では、こちらの手順で必要な環境をセットアップできます。

# -*- coding: utf-8 -*-
#
# k平均法による画像の減色処理
#
# 2015/04/24 ver1.0
#

import numpy as np
from numpy.random import randint
from PIL import Image

#------------#
# Parameters #
#------------#
Colors = [2, 3, 5, 16]  # 減色後の色数（任意の個数の色数を指定できます）


# k平均法による減色処理
def run_kmeans(pixels, k):
    cls = [0] * len(pixels)

    # 代表色の初期値をランダムに設定
    center = []
    for i in range(k):
        center.append(np.array([randint(256), randint(256), randint(256)]))
    print map(lambda x: x.tolist(), center)
    distortion = 0

    # 最大50回のIterationを実施
    for iter_num in range(50): 
        center_new = []
        for i in range(k):
            center_new.append(np.array([0,0,0]))
        num_points = [0] * k
        distortion_new = 0

        # E Phase: 各データが属するグループ（代表色）を計算
        for pix, point in enumerate(pixels):
            min_dist = 256*256*3
            point = np.array(point)
            for i in range(k):
                d = sum([x*x for x in point-center[i]])
                if d < min_dist:
                    min_dist = d
                    cls[pix] = i
            center_new[cls[pix]] += point
            num_points[cls[pix]] += 1
            distortion_new += min_dist

        # M Phase: 新しい代表色を計算
        for i in range(k):
            center_new[i] = center_new[i] / num_points[i]
        center = center_new
        print map(lambda x: x.tolist(), center)
        print "Distortion = %d" % distortion_new

        # Distortion(J)の変化が0.5%未満になったら終了
        if iter_num > 0 and distortion - distortion_new < distortion * 0.005:
            break
        distortion = distortion_new

    # 画像データの各ピクセルを代表色で置き換え
    for pix, point in enumerate(pixels):
        pixels[pix] = tuple(center[cls[pix]])

    return pixels
        
# Main
if __name__ == '__main__':
    for k in Colors:
        print "k=%d" % k
        # 画像ファイルの読み込み
        im = Image.open("photo.jpg")
        pixels = list(im.convert('RGB').getdata())
        # k平均法による減色処理
        result = run_kmeans(pixels, k)
        # 画像データの更新とファイル出力
        im.putdata(result) # Update image
        im.save("output%02d.bmp" % k, "BMP")