图片主色提取

提取主色

最近有个小任务，需要提取图片中的主色块。最基本的做法就是聚类找主色，但是目前需要自适应提取，也就是说需要对于不同的图片提取出不同的主色块数量。

1. RGB Or HSV

貌似大多数提取主色都是建立在RGB色彩空间中，但是就人眼感知而言对红色不太敏感而对蓝色较为敏感；而且，RGB色彩空间是利用三个颜色分量的线性组合来表示颜色相关性很高，所以RGB是一种不均匀的颜色空间。

再来看看HSV，对于单一颜色我们可以控制颜色角H，通过调整饱和度和明度就能得到同颜色系的不同颜色，这样的话我们聚类的结果也应该更为准确，所以我这里选择HSV颜色空间来进行聚类。

2. KMeans 自适应

我们都知道聚类效果的好坏有很多评价指标，比如SSE、轮廓系数……但是“肘部法”对于自适应来说并不适用，所以还是选择轮廓系数最大的k值作为最优k值。

并且对于图片聚类来说，一张图中的像素点很多，计算KMeans以及轮廓系数十分耗时，为了效率我进行了以下操作：

• 对输入图片进行缩放，缩小图片大小
• 使用MiniBatchKMeans代替``Kmeans`
• 计算轮廓系数时进行采样计算

3. 整体过程及代码

1. 读取图片，并缩小
2. 将RGB转为HSV进行聚类，找到最好的K值
3. 利用最好的K值再次聚类，得到最终结果并转回RGB
4. 可视化提取的主色块

import time
from functools import wraps
import matplotlib.pyplot as plt
import numpy as np
from skimage.color import rgb2hsv, hsv2rgb
from skimage import transform
from skimage.io import imread
from sklearn.cluster import MiniBatchKMeans
from sklearn.metrics import silhouette_score

# calculating time
def time_it(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        print(f"Spend time:{time.time() - start} s")
        return result

    return wrapper


@time_it
def get_best_k(src):
    K_list = [i for i in range(2, 8)]
    scores = []
    for i in K_list:
        mbk = MiniBatchKMeans(n_clusters=i, random_state=0)
        scores.append(silhouette_score(src, mbk.fit_predict(src), sample_size=int(src.shape[0] / 128)))
    index = scores.index(max(scores))
    best_k = K_list[index]
    print(f'best_k:', best_k)
    return best_k


def get_main_Color(src):
    # to_HSV
    img = rgb2hsv(src)
    h, w, d = img.shape
    img = np.reshape(img, (h * w, d))
    k = get_best_k(img)
    bk = MiniBatchKMeans(n_clusters=k, random_state=0)
    bk.fit(img)
    result = bk.cluster_centers_
    # to_RGB
    maincolor = hsv2rgb(result) * 255
    return maincolor


def maincolor_show(maincolor):
    N = len(maincolor)
    img = np.zeros((300, 100 * N, 3), np.uint8)

    for i in range(N):
        img[:, 100 * i:100 * (i + 1)] = [maincolor[i]]
    plt.imshow(img)
    plt.axis('off')
    plt.show()


def main():
    src = imread("/home/shelgi/图片/1.jpeg")
    plt.imshow(src)
    plt.axis('off')
    plt.show()

    src = transform.rescale(src, [0.4, 0.4, 1])
    maincolor = get_main_Color(src)
    print(maincolor)
    # visualization
    maincolor_show(maincolor)


if __name__ == '__main__':
    main()

4. 结果

后续

这个功能我觉得还是很可玩的，如果实时性有一定保障的话，我准备有空写个简单的服务，传一张图片显示提取主色块图，这样又水一篇blog。