机器学习让配色更简单

import cv2
import numpy as np
import matplotlib.pyplot as plt
import hypertools as hyp
import numpy as np
from sklearn.cluster import KMeans
from os import listdir
 
def get_leading_colors(directory):
    for filename in listdir(directory):
        img = cv2.imread(directory + filename)
 
        # by default, cv2 uses BGR so we need to change it to RGB
        img_data = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        r, g, b = cv2.split(img_data)
 
        # plot a t-SNE projection of each color channel
        hyp.plot(r, '.', reduce='TSNE', ndims=2, color='red',
                 title='2D t-SNE projection of the red channel',
                 size=[14, 8])
        hyp.plot(g, '.', reduce='TSNE', ndims=2, color='green',
                 title='2D t-SNE projection of the green channel',
                 size=[14, 8])
        hyp.plot(b, '.', reduce='TSNE', ndims=2, color='blue',
                 title='2D t-SNE projection of the blue channel',
                 size=[14, 8])
 
        # row * column, and number of color channels (3 because of RGB)
        img = img.reshape((img_data.shape[0] * img_data.shape[1], 3))
 
        # the number of clusters indicate how many leading colors we want
        model = KMeans(n_clusters=k, init='random', random_state=88)
        model.fit(img)
 
        hist = compute_histogram(model)
        rect = draw_leading_color_plot(hist, model.cluster_centers_)
 
        plt.axis('off')
        plt.imshow(rect)
        plt.show()

def compute_histogram(model):
    labels_list = np.arange(0, k + 1)
    # this histogram says how many pixels fall into one of the bins
    (hist, _) = np.histogram(model.labels_, bins=labels_list)
    hist = hist.astype('float')
    hist /= hist.sum()
 
    return hist

def draw_leading_color_plot(hist, centroids):
    # the first two values of np.zeros(...) represent the size of the rectangle
    # the 3 is because of RGB
    plot_width = 700
    plot_length = 150
    plot = np.zeros((plot_length, plot_width, 3), dtype='uint8')
    start = 0
 
    for (percent, color) in sorted(zip(hist, centroids), key=lambda x: x[0], reverse=True):
        end = start + (percent * plot_width)
        # append the leading colors to the rectangle
        cv2.rectangle(plot, (int(start), 0), (int(end), plot_length),
                      color.astype('uint8').tolist(), -1)
        print(color)
        start = end
 
    # return the rectangle chart
    return plot