Pythonで画像からメインカラーを抽出する

TL;DR

↑などの画像から、

↑みたいに、メインカラーをPythonで抽出します。

使うもの

• scikit-learn（k-means法）
• PIL、OpenCV（画像周り）

サンプルデータ

元々はDCGANモデル検証でのイラストの色変換などのためのデータセットでカラーヒントを用意するためだったのですが、ここではフリー素材の写真を使ってみます。（イラスト側でも問題なく対応できるのは確認済みです）

スクリプト

画像の読み込みと変換

花の画像を対象に進めていきます。

import PIL

from PIL import Image

import cv2

import sklearn

from sklearn.cluster import KMeans

cv2_img = cv2.imread('./flower.jpg')

OpenCV、読み込んだままだとカラーチャンネルの順番がRGBになっていないようなので変換をしておきます。

cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)

処理の都合、(縦, 横, カラーチャンネル)の形式の画像のテンソルを、(縦と横をフラット化させた値, カラーチャンネル)の2次元の形式に変換しておきます。

cv2_img.shape

(900, 1600, 3)

cv2_img = cv2_img.reshape(
    (cv2_img.shape[0] * cv2_img.shape[1], 3))

cv2_img.shape

(1440000, 3)

クラスタリング

k-means法でクラスタリングを行います。各ピクセルの色で、指定したクラスタ数でクラスタリングされ、その各クラスタの中央の座標の色をメインカラーとして扱います。

引数のn_clusterはクラスター数となります。今回はメインカラーを5色算出したいので、5を指定していますが、適宜調整してください。

cluster = KMeans(n_clusters=5)

対象画像に対してクラスタリングを行います。画像サイズによっては少し時間がかかります。

cluster.fit(X=cv2_img)

KMeans(algorithm='auto', copy_x=True, init='k-means++', max_iter=300,
    n_clusters=5, n_init=10, n_jobs=1, precompute_distances='auto',
    random_state=None, tol=0.0001, verbose=0)

cluster_centers_ という属性に、各クラスタの中央の値（RGBの値）が格納されます。

cluster.cluster_centers_

array([[211.60346671, 238.98895507, 231.41817764],
       [104.07336583, 143.33552179, 167.95419295],
       [183.88814356,  93.68673428,  49.43004867],
       [167.54667869, 212.4601815 , 211.04911242],
       [ 63.75784365,  60.18503292,  86.27467466]])

cluster.cluster_centers_.shape

(5, 3)

クラスタ数を5にし、RGB3チャンネルの画像なので、(5, 3)の配列となりました。

色を表示してみる

RGBの配列だとイメージがつきづらいので、PILを使って単色画像でそれぞれ表示してみます。

Pythonでは、02xと指定することで、0～255までの値を16進数の00～ffのフォーマットにしてくれるので、それを使って#ffffffといった文字列を用意してPILに渡します。

format(255, '02x')

'ff'

'#%02x%02x%02x' % (255, 255, 255)

'#ffffff'

浮動小数点数は不要なので、キャストしておきます。（16進数変換の都合）

cluster_centers_arr = cluster.cluster_centers_.astype(
    int, copy=False)

ループで回す都合、Jupyter上で全ての画像を表示するためにdisplay関数を使用します。（そのままだと最後の画像しか、アウトプットとして表示されないため）

from IPython.display import display

for rgb_arr in cluster_centers_arr:
    color_hex_str = '#%02x%02x%02x' % tuple(rgb_arr)
    color_img = Image.new(
        mode='RGB', size=(32, 32), color=color_hex_str)
    display(color_img)

元画像を表示してみて、比較してみます。

original_img = Image.open('./flower.jpg')

original_img.size

(1600, 900)

original_img.resize((320, 180))

いい感じです。

おまけで、PILを使って背景をグレーにし、横並びの画像を作ってみます。

# 幅と高さ64px × 横並び5画像 + 上下左右余白15pxずつの画像を作ります。
IMG_SIZE = 64
MARGIN = 15
width = IMG_SIZE * 5 + MARGIN * 2
height = IMG_SIZE + MARGIN * 2
print(width, height)

350 94

tiled_color_img = Image.new(
    mode='RGB', size=(width, height), color='#333333')

この段階では上記の用意した画像は、グレーの横長画像です。

この画像に、各クラスタの色の画像をペーストしていきます。PILのpaste関数を使います。引数のboxで、左の余白値、上の余白値といった具合でペーストする位置を設定できるので、ループのインデックスを加味して位置を指定します。

Pythonでループのインデックスを取りたい場合には、enumerate関数を使うとindex, valueといった形で扱えるのでシンプルです。

for i, rgb_arr in enumerate(cluster_centers_arr):
    color_hex_str = '#%02x%02x%02x' % tuple(rgb_arr)
    color_img = Image.new(
        mode='RGB', size=(IMG_SIZE, IMG_SIZE),
        color=color_hex_str)
    tiled_color_img.paste(
        im=color_img,
        box=(MARGIN + IMG_SIZE * i, MARGIN))

tiled_color_img

Adobe Color CCの探索画面みたいなカラーセット画像を用意することができました。

他の画像で試して見る

関数化して、他の写真でも試してみましょう。

def get_main_color_list_img(img_path):
    """
    対象の画像のメインカラーを算出し、色を横並びにしたPILの画像を取得する。
    
    Parameters
    ----------
    img_path : str
        対象の画像のパス。
    
    Returns
    -------
    tiled_color_img : Image
        色を横並びにしたPILの画像。
    """
    cv2_img = cv2.imread(img_path)
    cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
    cv2_img = cv2_img.reshape(
        (cv2_img.shape[0] * cv2_img.shape[1], 3))
    
    cluster = KMeans(n_clusters=5)
    cluster.fit(X=cv2_img)
    cluster_centers_arr = cluster.cluster_centers_.astype(
        int, copy=False)
    
    IMG_SIZE = 64
    MARGIN = 15
    width = IMG_SIZE * 5 + MARGIN * 2
    height = IMG_SIZE + MARGIN * 2
    
    tiled_color_img = Image.new(
        mode='RGB', size=(width, height), color='#333333')
    
    for i, rgb_arr in enumerate(cluster_centers_arr):
        color_hex_str = '#%02x%02x%02x' % tuple(rgb_arr)
        color_img = Image.new(
            mode='RGB', size=(IMG_SIZE, IMG_SIZE),
            color=color_hex_str)
        tiled_color_img.paste(
            im=color_img,
            box=(MARGIN + IMG_SIZE * i, MARGIN))
    return tiled_color_img

def get_original_small_img(img_path):
    """
    元画像の小さくリサイズしたPILの画像を取得する。
    
    Parameters
    ----------
    img_path : str
        対象の画像のパス。
    
    Returns
    -------
    img : Image
        リサイズ後の画像。
    """
    img = Image.open(fp=img_path)
    width = int(img.size[0] / 5)
    height = int(img.size[1] / 5)
    img = img.resize(size=(width, height))
    return img

get_original_small_img(img_path='./cat.jpg')

get_main_color_list_img(img_path='./cat.jpg')

get_original_small_img(img_path='./sky.jpg')

get_main_color_list_img(img_path='./sky.jpg')

get_original_small_img(img_path='./buildings.jpg')

get_main_color_list_img(img_path='./buildings.jpg')

実行環境

• Azure Notebooks

!python -V

Python 3.5.4 :: Anaconda custom (64-bit)

PIL.__version__

'4.3.0'

cv2.__version__

'3.3.1'

sklearn.__version__

'0.18.1'

参考サイト

OpenCV and Python K-Means Color Clustering

おまけ

ノートをgithubにアップしておきました。

20180503_画像からメインカラーを抽出する.ipynb