Spaces:
Running
Running
| #This file finds X dominant colors of a given image, using k-means and different channels of the image, | |
| #like RGB, and luminance. | |
| import colorsys | |
| import math | |
| import numpy as np | |
| from PIL import Image | |
| from sklearn.cluster import MiniBatchKMeans | |
| from collections import Counter | |
| PALETTE_WIDTH = 800 | |
| PALETTE_HEIGHT = 200 | |
| MAX_SIZE = 500,500 | |
| # def make_palette(file, n_clusters): | |
| # im = Image.open(file) | |
| # im.thumbnail(MAX_SIZE) | |
| # pixels = np.array([im.getpixel((x,y)) for x in range(0, im.size[0]) for y in range(0, im.size[1])]) | |
| # clt = MiniBatchKMeans(n_clusters = n_clusters, n_init = 1) | |
| # clt.fit(pixels) | |
| # return [[int(round(i)) for i in color] for color in clt.cluster_centers_] | |
| def make_palette(file, n_clusters): | |
| im = Image.open(file) | |
| im.thumbnail(MAX_SIZE) | |
| pixels = np.array([im.getpixel((x, y)) for x in range(im.size[0]) for y in range(im.size[1])]) | |
| clt = MiniBatchKMeans(n_clusters=n_clusters, n_init=1) | |
| clt.fit(pixels) | |
| cluster_centers = [[int(round(i)) for i in color] for color in clt.cluster_centers_] | |
| # sort cluster centers by number of pixels assigned to them | |
| counts = Counter(clt.labels_) | |
| sorted_counts = sorted(counts.items(), key=lambda x: -x[1]) | |
| cluster_centers = [cluster_centers[i] for i, _ in sorted_counts] | |
| return cluster_centers | |
| def perceived_brightness (r, g, b): | |
| """ | |
| Calculates perceived brightness for the given RGB values. | |
| code from Darel Rex Finley (http://alienryderflex.com/hsp.html) | |
| """ | |
| return math.sqrt((.299 * r * r) + (.587 * g * g) + (.114 * b * b)) | |
| def hsp_rank (r, g, b, mult=8): | |
| """ | |
| Combines hue, saturation, and perceived brightness info for a | |
| smoother sort. | |
| """ | |
| lum = perceived_brightness(r, g, b) | |
| h, s, v = colorsys.rgb_to_hsv(r, g, b) | |
| return (h * lum * mult, s * lum * mult, s * mult) | |
| def findDominantColors(imageFile, numberOfColors, showImage = True): | |
| src = imageFile | |
| n_clusters = numberOfColors | |
| # palette is ordered by its presence in the image | |
| cluster_centers = make_palette(src, n_clusters) | |
| # cluster_centers.sort(key=lambda rgb: hsp_rank(*rgb)) | |
| pixels = [] | |
| dominantColors = [] | |
| for i in range(n_clusters): | |
| color = tuple(cluster_centers[i]) | |
| dominantColors.append(color) | |
| # print("Color #{}: {}".format(i+1, color)) | |
| for j in range(PALETTE_WIDTH//n_clusters * PALETTE_HEIGHT): | |
| pixels.append(color) | |
| if showImage: | |
| img = Image.new('RGB',(PALETTE_HEIGHT,PALETTE_WIDTH)) | |
| img.putdata(pixels) | |
| img.show() | |
| return dominantColors | |
| def rgb_to_hex(rgb): | |
| return '#{:02x}{:02x}{:02x}'.format(rgb[0], rgb[1], rgb[2]) | |
| def hex_to_rgb(hex): | |
| hex = hex.lstrip('#') | |
| return tuple(int(hex[i:i+2], 16) for i in (0, 2, 4)) | |
| def create_color_palette(colors, palette_width=800, palette_height=200): | |
| """ | |
| Receives a list of colors in hex format and creates a palette image | |
| """ | |
| pixels = [] | |
| n_colors = len(colors) | |
| for i in range(n_colors): | |
| color = hex_to_rgb(colors[i]) | |
| for j in range(palette_width//n_colors * palette_height): | |
| pixels.append(color) | |
| img = Image.new('RGB', (palette_height, palette_width)) | |
| img.putdata(pixels) | |
| img.show() | |
| ##############pipeline for testing | |
| # print("Dominant colors:") | |
| # dominantColors=findDominantColors('./estampa-test.png', 10, False) | |
| # dominantColorHex=[rgb_to_hex(color) for color in dominantColors] | |
| # print(dominantColorHex) | |
| # # ['#d54b18', '#555015', '#98680e', '#1f3919', '#da8007', '#a50637', '#b50406', '#d52350', '#e2090f', '#f8545c'] |