Python 提取图片主色调

效果

有个要提取图片主色调的需求，记录一下。

代码编写

import numpy as np
import cv2
from sklearn.cluster import KMeans
from skimage.color import rgb2lab, deltaE_cie76
from collections import Counter# 创建默认变量# 这是用于确定背景的阈值
remove_background_threshold = 235
# 这是用于确定要提取的颜色数量的参数
extract_colors_num_colors = 8
# 这是用于确定两种颜色是否相似的阈值
is_similar_color_threshold = 30def remove_background(image, threshold=remove_background_threshold):# Convert image to grayscale# 将图像转换为灰度图gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)# Create a binary mask of the background# 创建一个背景的二值掩码_, mask = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY_INV)# Apply the mask to the original image# 将掩码应用到原始图像上result = cv2.bitwise_and(image, image, mask=mask)return resultdef get_color_name(lab):colors = {"red": [53.23, 80.11, 67.22],"green": [87.74, -86.18, 83.18],"blue": [32.30, 79.19, -107.86],"yellow": [97.14, -21.55, 94.48],"cyan": [91.11, -48.09, -14.13],"magenta": [60.32, 98.24, -60.83],"black": [0, 0, 0],"white": [100, 0.00, -0.01],"gray": [53.59, 0, 0],"orange": [67.79, 43.30, 74.93],"purple": [29.78, 58.94, -36.50],"brown": [38.91, 19.36, 22.29],"pink": [88.22, 17.75, 3.18],"dark red": [39.35, 62.75, 49.91],"light blue": [79.19, -11.03, -26.23],"dark green": [35.49, -46.47, 35.45],"light green": [88.72, -42.89, 57.40],"navy": [16.73, 37.09, -65.49],"burgundy": [28.71, 49.35, 26.27],"beige": [89.02, -1.39, 11.09],"olive": [51.87, -12.93, 56.67],"teal": [49.31, -28.83, -8.48],"maroon": [25.64, 45.52, 20.77],"forest green": [36.23, -37.96, 30.20],}min_dist = float('inf')closest_color = Nonefor color_name, color_lab in colors.items():dist = np.sqrt(sum((lab - color_lab) ** 2))if dist < min_dist:min_dist = distclosest_color = color_namereturn closest_colordef is_similar_color(color1, color2, threshold=is_similar_color_threshold):lab1 = rgb2lab([[color1]])[0][0]lab2 = rgb2lab([[color2]])[0][0]diff = deltaE_cie76(lab1, lab2)return diff < thresholddef extract_colors(image_path, num_colors=extract_colors_num_colors):# Read image# 读取图像image = cv2.imread(image_path)image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)# Remove background# 去除背景image_no_bg = remove_background(image)# Reshape image# 重塑图像pixels = image_no_bg.reshape(-1, 3)# Remove black pixels (background)# 删除黑色像素（背景）pixels = pixels[~np.all(pixels == [0, 0, 0], axis=1)]if len(pixels) == 0:return []  # Return empty list if all pixels were removed   返回空列表，如果所有像素都被删除# Perform k-means clustering# 执行k-means聚类kmeans = KMeans(n_clusters=num_colors, random_state=42, n_init=10)kmeans.fit(pixels)# Get the colors# 获取颜色colors = kmeans.cluster_centers_# Convert colors to LAB space# 将颜色转换为LAB空间colors_lab = rgb2lab(colors.reshape(1, -1, 3)).reshape(-1, 3)# Get color names# 获取颜色名称color_names = [get_color_name(color) for color in colors_lab]# Count pixels for each cluster# 计算每个聚类的像素数pixel_counts = Counter(kmeans.labels_)# Sort colors by frequency# 按频率对颜色进行排序sorted_colors = sorted(zip(color_names, colors, pixel_counts.values()), key=lambda x: x[2], reverse=True)# Remove similar colors# 删除相似颜色unique_colors = []for color_name, color_rgb, count in sorted_colors:if color_name not in [uc[0] for uc in unique_colors] and \not any(is_similar_color(color_rgb, existing_color) for _, existing_color, _ in unique_colors):unique_colors.append((color_name, color_rgb, count))if len(unique_colors) == 3:breakreturn [color_name for color_name, _, _ in unique_colors]# Example usage
image_path = r'D:\Project\Python_Project\yitiaolong\expimage_1.jpg'
main_colors = extract_colors(image_path)
print("Main colors:", main_colors)