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 = 30

def 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 result

def 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 = None
    for color_name, color_lab in colors.items():
        dist = np.sqrt(sum((lab - color_lab) ** 2))
        if dist < min_dist:
            min_dist = dist
            closest_color = color_name
    return closest_color

def 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 < threshold

def 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:
            break
    
    return [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)