使用Python提取图片中的主体颜色

上次国庆节去谷歌开发者大会,体验了很多有趣的人工智能项目。

其中有一个颜色匹配的环节,叫做“AI调色板,解码缤纷艺术世界”,让我觉得很有意思,回来后,我计划自己实现一个类似的功能。

简单点讲,原理其实很简单,

  • 给系统输入一张图片
  • 系统检测输入图片中的什么色彩占比比较大
  • 根据上一步识别的色彩,到庞大的图片库中匹配与之类似的图片

首先,我们需要想办法提取出图片中的主体颜色。这里我用了传统的机器学习算法——k-means来对图片上的像素值进行聚类。

from skimage import io
from sklearn.cluster import KMeans
import numpy as np
import warnings

warnings.filterwarnings('ignore')

img_file = 'city.jpg'

# k-means中的k值,即选择几个中心点
k = 5

# 读图片
img = io.imread(img_file)
# 转换数据维度
img_ori_shape = img.shape
img1 = img.reshape((img_ori_shape[0] * img_ori_shape[1], img_ori_shape[2]))
img_shape = img1.shape

# 获取图片色彩层数
n_channels = img_shape[1]

estimator = KMeans(n_clusters=k, max_iter=4000, init='k-means++', n_init=50)  # 构造聚类器
estimator.fit(img1)  # 聚类
centroids = estimator.cluster_centers_  # 获取聚类中心

# 使用算法跑出的中心点,生成一个矩阵,为数据可视化做准备
result = []
result_width = 200
result_height_per_center = 80
for center_index in range(k):
    result.append(np.full((result_width * result_height_per_center, n_channels), centroids[center_index], dtype=int))
result = np.array(result)
result = result.reshape((result_height_per_center * k, result_width, n_channels))

# 保存图片
io.imsave(img_file + '.result.bmp', result)

程序里的city.jpg就是下面这张图片

代码里,我们设置的k值是5,下面程序中计算出来的中心点,每个点由三个点构成,分别对应R,G,B通道。

  1. [41.1908632, 42.93972933, 42.17768105]
  2. [235.53604525, 229.25576283, 230.17416456]
  3. [196.38345078, 162.31462071, 129.28781508]
  4. [144.83752455, 113.68305275,  87.47028264]
  5. [89.70849154, 71.46176502, 58.98393717]

将这些数据进行可视化,得到下图。可以看出,分类的颜色基本上是city.jpg中占比比较大的颜色。

我们多运行几次试试看返回的颜色是不是会变化。

运行8次堆叠的结果
这里有可能是k-means训练过程的随机性导致的,但上面的颜色中,颜色主要是5类,只是顺序稍微有所改变。

下次有时间把接下来的内容进行完善:

  • 用python实现k-means,不依赖sklearn
  • 将多张图片的主体颜色结果进行匹配,对相似颜色的图片进行聚类

此项目代码已经上传到github上了:https://github.com/shulisiyuan/mainColor

The End


已发布

分类

标签:

评论

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注