Pynote

Python、機械学習、画像処理について

OpenCV - k 平均法 (k-means) を使い、画像の代表色を取得する方法

概要

OpenCV で k 平均法 (k-means) を使い、画像の代表色を取得する方法を紹介する。
以下の用途で利用できる。

  • 画像で使われている主な色を取得する。
  • 画像を見た目にあまり影響がないように減色する。(例: 1000色で表されている画像を8色で表現する。)

cv2.kmeans

retval, bestLabels, centers = cv2.kmeans(
    data, K, bestLabels, criteria, attempts, flags[, centers])
  • 引数
  • 返り値:
    • retval: compactness measure の値。
    • bestLabels: 各サンプルに割り当てられたクラスタ
    • centers: クラスタの中心の一覧

criteria

criteria は最大反復回数及び移動量の閾値の2通りを指定できる。

attempts

最初、クラスタの中心は適当な点で初期化されるので、アルゴリズムの結果は毎回異なる。
attempts に2以上の値を指定した場合、アルゴリズムをその回数分実行し、compactness measure の値が最小となる結果、つまり、最良の結果を返す。

retval

サンプルとそのサンプルが属するクラスタの中心との2乗距離のこと。

l_i = \texttt{labels}_i としたとき、
\displaystyle
\texttt{retval} = \sum_i \| \texttt{samples}_i - \texttt{centers}_{l_i} \|^2

サンプルコード

入力画像

k 平均法で代表色を計算する。

import cv2
import numpy as np
import matplotlib.pyplot as plt


# 画像を読み込む。
img = cv2.imread("sample.jpg")
print(img.shape)  # (358, 600, 3)

# 画像で使用されている色一覧。(W * H, 3) の numpy 配列。
colors = img.reshape(-1, 3)
# cv2.kmeans に渡すデータは float 型である必要があるため、キャストする。
colors = colors.astype(np.float32)

# k平均法でクラスタリングする。

# クラスタ数
K = 5

# 最大反復回数: 10、移動量の閾値: 1.0
criteria = cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 10, 1.0

ret, label, center = cv2.kmeans(
    colors, K, None, criteria, attempts=10, flags=cv2.KMEANS_RANDOM_CENTERS
)

print(f"ret: {ret}, label: {label.shape}, center: {center.shape}")
# ret: 127443.79220199585, label: (48380, 1), center: (8, 3)

label は各サンプルが割り当てられたクラスタを表す (NumSamples, 1) の numpy 配列である。
center はクラスタの中心点を表す (K, M) の numpy 配列である。

例えば、i 番目のサンプルが割り当てられたクラスタの中心点を取得したい場合は、center[label[i]] とすればよい。

# 各クラスタに属するサンプル数を計算する。
_, counts = np.unique(label, axis=0, return_counts=True)

# matplotlib で棒グラフを作成する。
fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(10, 3))
fig.subplots_adjust(wspace=0.5)

# matplotlib の引数の仕様上、[0, 1] にして、(R, G, B) の順番にする。
bar_color = [(r / 255, g / 255, b / 255) for b, g, r in center]
bar_text = [f"({r}, {g}, {b})" for b, g, r in center.astype(np.uint8)]

# 画像を表示する。
ax1.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
ax1.set_axis_off()

# ヒストグラムを表示する。
ax2.barh(np.arange(K), counts, color=bar_color, tick_label=bar_text)
plt.show()


画像の各画素をクラスタの色で置き換える。(減色処理)

# 各画素を k平均法の結果に置き換える。
dst = center[label.ravel()].reshape(img.shape)
dst = dst.astype(np.uint8)

fig, ax = plt.subplots()
ax.imshow(cv2.cvtColor(dst, cv2.COLOR_BGR2RGB))
ax.set_axis_off()

plt.show()

画像で使われている色を5色に減らしたが、元の画像が想像できる程度には色合いが保たれている。