Skip to content

Instantly share code, notes, and snippets.

Last active October 10, 2023 13:45
Show Gist options
  • Save simonmesmith/a1c3fdef3d8e9a03cd170cc3e7a5a596 to your computer and use it in GitHub Desktop.
Save simonmesmith/a1c3fdef3d8e9a03cd170cc3e7a5a596 to your computer and use it in GitHub Desktop.


Extract colors from images and attach them as swatches. See an example here.

Takes one argument: number_of_color_clusters (int), which is pretty self-explanatory. No? Okay, it means how many color clusters (e.g. reddish colors, bluish colors, etc.) to return and add to the swatches.


input_path = "input_image.jpg"
output_path = "output_image.jpg"
colorgrammed_image = create(image_path=input_path, number_of_color_clusters=10)
cv2.imwrite(output_path, colorgrammed_image)

PS: I owe a debt to various places where I learned how to do this and copied some code snippets. I don't remember them all but now wish I did. If you recognize this as something you've worked on, and want credit, please let me know as I'm happy to give it!

import cv2
import numpy as np
def create(image_path: str, number_of_color_clusters: int) -> list:
# Read in the image.
img = cv2.imread(image_path)
# Get the image's height, width, and channels.
img_height, img_width, img_channels = np.shape(img)
# Convert the image into a list of pixels, with each pixel representing a color.
pixels = np.reshape(img, (img_height * img_width, img_channels))
# Turn each number in the pixel array into a float32, needed for the kmeans() function.
pixels = np.float32(pixels)
# Specify kmeans clustering settings.
kmeans_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0) # How clustering should work
kmeans_flags = cv2.KMEANS_RANDOM_CENTERS # Start clustering with a random set of samples
# Cluster colors and set pixel_labels to the cluster label for each pixel, and color_clusters to the colors we've clustered pixels into.
_, pixel_labels, color_clusters = cv2.kmeans(pixels, number_of_color_clusters, None, kmeans_criteria, 10, kmeans_flags)
# Calculate the percentage of pixels that each color cluster represents.
(pixel_label_counts, _) = np.histogram(pixel_labels, bins=number_of_color_clusters)
pixel_label_percentages = pixel_label_counts / pixel_label_counts.sum()
# Create a canvas on which to display the colors in percentages.
color_bar_width = img_width # Because we stack the original image on top of the color bar, so they need to be the same height
color_bar_height = int(img_height * .2) # Make the height proportional
color_bar = np.zeros((color_bar_height, color_bar_width, 3), dtype = "uint8")
# Merge the colours and percentages they represent. Then order by percentage in descending order so more prominent color clusters come first.
colors_and_percentages = list(zip(color_clusters, pixel_label_percentages))
colors_and_percentages.sort(key=lambda tup: tup[1], reverse=True)
# Create a variable to represent the starting point of each color in the bar.
start_x = 0
# For each color and percent...
for (color, percent) in colors_and_percentages:
# Plot the width of the color in the bar.
end_x = start_x + (percent * color_bar_width)
# Create the rectangle with OpenCV.
cv2.rectangle(color_bar, (int(start_x), 0), (int(end_x), color_bar_height), color.astype("uint8").tolist(), -1)
# Update the starting point for the next colour.
start_x = end_x
# Merge the color bar and original image.
combined = np.zeros((img_height + color_bar_height, img_width, 3), np.uint8)
combined[:img_height, :img_width, :3] = img
combined[img_height:img_height+color_bar_height, :img_width, :3] = color_bar
# Return the new image.
return combined
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment