extract-colors.py

from sklearn.cluster import KMeans
from collections import Counter
import numpy as np
import cv2

def get_image(pil_image):
    nimg = np.array(pil_image)
    image = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

def get_labels(rimg):
    clf = KMeans(n_clusters=5)
    labels = clf.fit_predict(rimg)
    return labels, clf

def get_closest_color(colors):
    white = (255, 255, 255)
    closest_color = min(colors, key=lambda c: np.linalg.norm(np.array(c) - white))
    return closest_color

def RGB2HEX(color):
    return "#{:02x}{:02x}{:02x}".format(int(color[0]), int(color[1]), int(color[2]))

def extract_colors_and_closest_to_white(image_path):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    reshaped_img = img.reshape(img.shape[0] * img.shape[1], img.shape[2])
    labels, clf = get_labels(reshaped_img)
    counts = Counter(labels)
    center_colors = clf.cluster_centers_
    ordered_colors = [center_colors[i] for i in counts.keys()]
    hex_colors = [RGB2HEX(ordered_colors[i]) for i in counts.keys()]

    closest_color_to_white = get_closest_color(center_colors)
    hex_closest_color_to_white = RGB2HEX(closest_color_to_white)
    return hex_colors, hex_closest_color_to_white