使用 Python 探测图像的主色调

记录使用 Python 探测一张图像/图片的主色调。

环境与依赖

环境 Windows 10 + Python 3.x;

需安装的 Python 包

cv2, numpy, matplotlib, sklearn

准备

创建一个目录，作为工作目录，在下面放几张非纯色的图片。

代码

在刚才创建的工作目录下，创建文件 dmc.py，粘贴如下内容：

# -*- coding: utf-8 -*-
# @Author: suifengtec
# @Date:   2017-06-18 15:20:24
# @Last Modified by:    suifengtec
# @Last Modified time: 2018-10-08 13:27:34

# 探测图片的主色调
# Detecting Dominant Color on an Image
#
# python dmc.py
#
#
import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import os
import sys


class DetectImageDominantColors:
    def __init__(self, filePath):
        self.filePath = filePath
        self.dominantColorsCluster = self.getClusters()
        self.dominantColorsPercentageArray = self.getColorsPercentageArr(
            self.dominantColorsCluster)

    def getClusters(self):
        """

        """
        if os.path.exists(self.filePath) == False:
            print('路径不存在:', self.filePath)
            return np.array([])
        if os.path.isfile(self.filePath) == False:
            print('这不是个文件:', self.filePath)
            return np.array([])

        try:
            srcBGR = cv2.imread(self.filePath)
        except OSError:
            print('OSError: 未能打开文件', self.filePath)
        except AssertionError:
            print('AssertionError: 未能打开文件', self.filePath)
        except cv2.error as e:
            print('AssertionError: 未能打开文件', str(e), file=sys.stderr)
        except:
            print('未知错误: 未能打开文件', self.filePath)

        # RGB
        # rgb = bgr[...,::-1]
        # gbr = rgb[...,[2,0,1]]
        #
        destRGB = cv2.cvtColor(srcBGR, cv2.COLOR_BGR2RGB)
        # 表示为行*列, 通道号
        destRGB = destRGB.reshape((destRGB.shape[0] * destRGB.shape[1], 3))
        # cluster number
        clstr = KMeans(n_clusters=3)

        clstr.fit(destRGB)
        return clstr

    def getColorsPercentageArr(self, cltsr):
        """
        使用 k clusters 创建一个直方图
        Arguments:
            cltsr {sklearn.cluster.k_means_.KMeans}
            实例
            KMeans(algorithm='auto', copy_x=True, init='k-means++', max_iter=300,
        n_clusters=3, n_init=10, n_jobs=None, precompute_distances='auto',
        random_state=None, tol=0.0001, verbose=0)

        Returns:
            pArray {numpy.ndarray} - - 各种主色的所占百分比的一维向量/数组
            实例
            [0.42676364 0.27444444 0.29879192]
        """

        numLabels = np.arange(0, len(np.unique(cltsr.labels_)) + 1)
        (pArray, _) = np.histogram(cltsr.labels_, bins=numLabels)
        pArray = pArray.astype("float")
        pArray /= pArray.sum()
        # pArray.sort(axis=0)
        return pArray

    def run(self, outputBar=False):
        """
        据记录和图心输出图的主色调图
        Arguments:
            hist {numpy.ndarray} - - 色标记录
            centroids {numpy.ndarray} - - 图心
        Returns:
            bar {numpy.ndarray} - - 直方图的数据
        """
        dominantColorsPercentageArray = self.dominantColorsPercentageArray
        if dominantColorsPercentageArray.size == 0:
            return False

        centroids = self.dominantColorsCluster.cluster_centers_

        bar = np.zeros((50, 300, 3), dtype="uint8")
        startX = 0
        for (percent, color) in zip(dominantColorsPercentageArray, centroids):
            # 绘制每个cluster的相对百分比
            endX = startX + (percent * 300)
            cv2.rectangle(bar, (int(startX), 0), (int(endX), 50),
                          color.astype("uint8").tolist(), -1)
            startX = endX
        if outputBar == True:
            plt.axis("off")
            plt.imshow(bar)
            plt.show()
        return bar


"""
测试
"""


def main(args):

    if len(args) == 1:

        file = os.path.join(os.getcwd(),  args[0])
        app = DetectImageDominantColors(file)
        bar = app.run(outputBar=True)
        # print(bar)
    else:
        print("使用示例:n python.dmc.py a.jpg")


if __name__ == '__main__':
    main(sys.argv[1:])

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

# -*- coding: utf-8 -*-

# @Author: suifengtec

# @Date: 2017-06-18 15:20:24

# @Last Modified by: suifengtec

# @Last Modified time: 2018-10-08 13:27:34

# 探测图片的主色调

# Detecting Dominant Color on an Image

# python dmc.py

import cv2

import numpy as np

import matplotlib.pyplot as plt

from sklearn.cluster import KMeans

import os

import sys

class DetectImageDominantColors:

def __init__(self, filePath):

self.filePath = filePath

self.dominantColorsCluster = self.getClusters()

self.dominantColorsPercentageArray = self.getColorsPercentageArr(

self.dominantColorsCluster)

def getClusters(self):

"""

if os.path.exists(self.filePath) == False:

print('路径不存在:', self.filePath)

return np.array([])

if os.path.isfile(self.filePath) == False:

print('这不是个文件:', self.filePath)

return np.array([])

try:

srcBGR = cv2.imread(self.filePath)

except OSError:

print('OSError: 未能打开文件', self.filePath)

except AssertionError:

print('AssertionError: 未能打开文件', self.filePath)

except cv2.error as e:

print('AssertionError: 未能打开文件', str(e), file=sys.stderr)

except:

print('未知错误: 未能打开文件', self.filePath)

# RGB

# rgb = bgr[...,::-1]

# gbr = rgb[...,[2,0,1]]

destRGB = cv2.cvtColor(srcBGR, cv2.COLOR_BGR2RGB)

# 表示为行*列, 通道号

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

# cluster number

clstr = KMeans(n_clusters=3)

clstr.fit(destRGB)

return clstr

def getColorsPercentageArr(self, cltsr):

"""

使用 k clusters 创建一个直方图

Arguments:

cltsr {sklearn.cluster.k_means_.KMeans}

实例

KMeans(algorithm='auto', copy_x=True, init='k-means++', max_iter=300,

n_clusters=3, n_init=10, n_jobs=None, precompute_distances='auto',

random_state=None, tol=0.0001, verbose=0)

Returns:

pArray {numpy.ndarray} - - 各种主色的所占百分比的一维向量/数组

实例

[0.42676364 0.27444444 0.29879192]

"""

numLabels = np.arange(0, len(np.unique(cltsr.labels_)) + 1)

(pArray, _) = np.histogram(cltsr.labels_, bins=numLabels)

pArray = pArray.astype("float")

pArray /= pArray.sum()

# pArray.sort(axis=0)

return pArray

def run(self, outputBar=False):

"""

据记录和图心输出图的主色调图

Arguments:

hist {numpy.ndarray} - - 色标记录

centroids {numpy.ndarray} - - 图心

Returns:

bar {numpy.ndarray} - - 直方图的数据

"""

dominantColorsPercentageArray = self.dominantColorsPercentageArray

if dominantColorsPercentageArray.size == 0:

return False

centroids = self.dominantColorsCluster.cluster_centers_

bar = np.zeros((50, 300, 3), dtype="uint8")

startX = 0

for (percent, color) in zip(dominantColorsPercentageArray, centroids):

# 绘制每个cluster的相对百分比

endX = startX + (percent * 300)

cv2.rectangle(bar, (int(startX), 0), (int(endX), 50),

color.astype("uint8").tolist(), -1)

startX = endX

if outputBar == True:

plt.axis("off")

plt.imshow(bar)

plt.show()

return bar

"""

测试

"""

def main(args):

if len(args) == 1:

file = os.path.join(os.getcwd(), args[0])

app = DetectImageDominantColors(file)

bar = app.run(outputBar=True)

# print(bar)

else:

print("使用示例:n python.dmc.py a.jpg")

if __name__ == '__main__':

main(sys.argv[1:])

测试

我找了一张图，是维基百科的：

https://en.wikipedia.org/wiki/Spartan_(apple)#/media/File:Spartan_apple.jpg

1	https://en.wikipedia.org/wiki/Spartan_(apple)#/media/File:Spartan_apple.jpg

apple

下载到本地后,重命名为 a.jpg, 在命令行工具中：

python dmc.py a.jpg

1	python dmc.py a.jpg

可看到：

detecting dominant colors on an image using python

不理想的地方是：比较慢，如果用于Web,需要和任务队列配合使用。

环境与依赖

准备

代码

测试

相关文章