機械学習を応用して画像のポスタリゼーション

K-means法は教師なし学習の中でクラスタリングを行いますが、そのアルゴリズムの中でクラスタの重心を求めます。
ところで画像は配列にするとRGBが縦×横になっているわけですが、そのRGBのデータを取り出すことでデータセットにしてクラスタの数を8個にして各色の重心の色を使います。

※今回はGoogle Colabを使っているのでインデントが少し変ですがおおめに見てください。

ライブラリのインポート

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

画像の読み込み

img = cv2.imread("/content/drive/MyDrive/lena_square.png")
plt.imshow(img)
plt.show()

なんかRGBがおかしいのはGoogle Colabの仕様なのかOpenCVの仕様なのか

RGBのデータセットを作る

arr = np.array(img).tolist()
data = []
for i in range(len(arr)):
  for j in range(len(arr[i])):
    data.append(arr[i][j])

これでRGBのテーブルデータが作れました。

モデルの作成と重心の確認

model = KMeans(n_clusters=8)
model.fit(data)
xyz = model.cluster_centers_
print(xyz)

[[172.18651769 194.6938175  231.58912477]
 [137.42536743 102.83654745 156.32401665]
 [ 65.31620059  23.47032448  93.80382301]
 [ 81.58072177  69.04411444 172.72797008]
 [135.19572287 156.5699912  225.06082865]
 [ 81.86449534  49.36292781 130.18762519]
 [ 95.78363064  96.39556654 205.78326798]
 [116.89396007 127.20808313 213.13809171]]

ポスタリゼーションした時にこの色が使われます。

色を重心に変換

arr2 = []
for i in range(len(arr)):
  tmp2 = []
  for j in range(len(arr[i])):
    tmp = []
    for k in range(len(xyz)):
      dst = np.sqrt((xyz[k][0]-arr[i][j][0])**2+(xyz[k][1]-arr[i][j][1])**2+(xyz[k][2]-arr[i][j][2])**2)
      tmp.append(dst)
    tmp2.append(xyz[np.argmin(tmp)])
  arr2.append(tmp2)

8色の中で一番ユークリッド距離が近い色を採用します。
これで8色になりました。

画像を確認

plt.imshow(np.array(arr2).astype("int"))
plt.show()

こんな感じになりました。割と自然にポスタリゼーションできたんじゃないかと思います。
比較用↓

まとめ

機械学習でこんなこともできるんですよ