Find the dominant colors in an image - fast!

cat.jpg

find_dominant_colors.py

from skimage import io
from operator import itemgetter
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cosine

# number of colors to fetch
top_colors = 6

# get an image matrix
#img = np.random.randint(2**8, size=(10,3))
url = 'https://static1.squarespace.com/static/55c945e0e4b04386fb9f8162/55eb967ee4b0296394987fcf/596cf08d4c0dbf6dda54f2f2/1500311745866/cat-111793.jpg'
img = io.imread(url).reshape(-1, 3)

img &= 0b11111111

# pluck out the r, g, b channels
r = img[:,0]
g = img[:,1]
b = img[:,2]

# specify the number of bins & range of values in each dimension
nbins = 12
_bins = (nbins,nbins,nbins)
_range = ((0,2**8),(0,2**8),(0,2**8))

# count observations in each bin range combination
H, [br, bg, bb]=np.histogramdd((r,g,b), bins=_bins, range=_range)
ir, ig, ib = np.where(H)
bin_counts = list( zip(br[ir], bg[ig], bb[ib], H[ir,ig,ib]) )
    
# sort the bins to find the most common bins
sorted_bins = sorted(bin_counts, key=itemgetter(3), reverse=True)

# select the n top bins, but make sure none are too similar
selected = []
selecting = True
while selecting:
  for i in sorted_bins:
    if not any([cosine(i[:-1], j) < .01 for j in selected]):
      selected.append(i[:-1])
    if len(selected) == top_colors:
      selecting = False
      break

# paint these colors - matplot expects them as {0:1}
palette = np.array(selected).astype(float) / 255
plt.imshow([palette])
plt.show()

top_colors.png