14-7. k-meansで画像の階調を下げる

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans


#============================
# パラメータ
#============================
#画像
file = 'hoge.jpg'

#何色に落とすか
n_clusters = 5


#============================
# 前処理
#============================
#画像読み込み
img = Image.open(file).convert('RGB'); img.close

#画像をたて200pxにリサイズする
w, h = int(img.size[0] * 200 / img.size[1]), 200
img_resize = img.resize((w, h), Image.BICUBIC)

#表示
print('img_resize')
plt.imshow(img_resize)
plt.show()

#(-1, 3)次元の配列に変形、-1は適当に判断してくれる
img_array = np.array(img_resize).reshape(-1, 3)
print('img_array.shape\n{}'.format(img_array.shape))


#============================
# k-means
#============================
#k-meansのオブジェクト
km = KMeans(n_clusters=n_clusters)

#計算を実行
km.fit(img_array)

#各クラスタの重心
km_centers = km.cluster_centers_
print('km_centers\n{}'.format(km_centers))

#計算結果に応じてクラスタリング（クラスタ番号の配列が返ってくる）
km_predict = km.predict(img_array)
print('km_predict.shape\n{}'.format(km_predict.shape))


#============================
# 後処理
#============================
#クラスタ番号を２次元に戻す
img_predict = km_predict.reshape(h, w)
print('img_predict.shape\n{}'.format(img_predict.shape))

#表示
print('img_predict')
plt.imshow(img_predict)
plt.show()

#クラスタ番号をクラスタ重心に置き換える、(h, w, 3)次元になる
img_final = np.empty((h, w, 3), dtype=int)
for i in range(n_clusters):
    img_final[img_predict==i] = km_centers[i]
print('img_final.shape\n{}'.format(img_final.shape))

#表示
print('img_final')
plt.imshow(img_final)
plt.show()