vlambda博客
学习文章列表

import matplotlib.pyplot as pltfrom mpl_toolkits.mplot3d import Axes3Dimport matplotlib.colors as color
import numpy as npfrom collections import Counterimport pandas as pdimport math
from sklearn.cluster import KMeansfrom PIL import Imageimport webcolorsimport json
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--clusters", help="No. of clusters")parser.add_argument("--imagepath", help="Path to input image")
args = parser.parse_args()
IMG_PATH = args.imagepath if args.imagepath else "images/poster.jpg"CLUSTERS = args.clusters if args.clusters else 5
WIDTH = 128HEIGHT = 128

with open('colors.json') as clr: color_dict = json.load(clr)

def TrainKMeans(img): new_width, new_height = calculate_new_size(img) image = img.resize((new_width, new_height), Image.ANTIALIAS) img_array = np.array(image) img_vector = img_array.reshape((img_array.shape[0] * img_array.shape[1], 3)) ''' ---------- Training K-Means Clustering Algorithm ---------- ''' kmeans = KMeans(n_clusters = CLUSTERS, random_state=0) labels = kmeans.fit_predict(img_vector)  hex_colors = [rgb_to_hex(center) for center in kmeans.cluster_centers_] color_name = {} for c in kmeans.cluster_centers_: h, name = findColorName(c) color_name[h] = name  img_cor = [[*x] for x in img_vector] ''' img_cor is a nested list of all the coordinates (pixel -- RGB value) present in the image ''' cluster_map = pd.DataFrame() cluster_map['position'] = img_cor cluster_map['cluster'] = kmeans.labels_ cluster_map['x'] = [x[0] for x in cluster_map['position']] cluster_map['y'] = [x[1] for x in cluster_map['position']] cluster_map['z'] = [x[2] for x in cluster_map['position']] cluster_map['color'] = [hex_colors[x] for x in cluster_map['cluster']] cluster_map['color_name'] = [color_name[x] for x in cluster_map['color']] print(cluster_map) return cluster_map, kmeans

def calculate_new_size(image): ''' We are resizing the image (one of the dimensions) to 128 px and then, scaling the other dimension with same height by width ratio. ''' if image.width >= image.height: wperc = (WIDTH / float(image.width)) hsize = int((float(image.height) * float(wperc))) new_width, new_height = WIDTH, hsize else: hperc = (HEIGHT / float(image.height)) wsize = int((float(image.width) * float(hperc))) new_width, new_height = wsize, HEIGHT return new_width, new_height

def rgb_to_hex(rgb): ''' Converting our rgb value to hex code. ''' hex = color.to_hex([int(rgb[0])/255, int(rgb[1])/255, int(rgb[2])/255]) print(hex)  return hex

def findColorName(rgb): ''' Finding color name :: returning hex code and nearest/actual color name ''' aname, cname = get_colour_name((int(rgb[0]), int(rgb[1]), int(rgb[2]))) hex = color.to_hex([int(rgb[0])/255, int(rgb[1])/255, int(rgb[2])/255]) if aname is None: name = cname else: name = aname return hex, name  def closest_colour(requested_colour): ''' We are basically calculating euclidean distance between our set of RGB values with all the RGB values that are present in our JSON. After that, we are looking  at the combination RGB (from JSON) that is at least distance from input RGB values, hence finding the closest color name. ''' min_colors = {} for key, name in color_dict['color_names'].items(): r_c, g_c, b_c = webcolors.hex_to_rgb("#"+key) rd = (r_c - requested_colour[0]) ** 2 gd = (g_c - requested_colour[1]) ** 2 bd = (b_c - requested_colour[2]) ** 2 min_colors[math.sqrt(rd + gd + bd)] = name #print(min(min_colours.keys())) return min_colors[min(min_colors.keys())]
def get_colour_name(requested_colour): ''' In this function, we are converting our RGB set to color name using a third  party module "webcolors".  RGB set -> Hex Code -> Color Name  By default, it looks in CSS3 colors list (which is the best). If it cannot find hex code in CSS3 colors list, it raises a ValueError which we are handling using our own function in which we are finding the closest color to the input RGB set. ''' try: closest_name = actual_name = webcolors.rgb_to_name(requested_colour) except ValueError: closest_name = closest_colour(requested_colour) actual_name = None return actual_name, closest_name

def plotColorClusters(img): cluster_map, kmeans = TrainKMeans(img) fig = plt.figure() ax = Axes3D(fig)  # grouping the data by color hex code and color name to find the total count of # pixels (data points) in a particular cluster mydf = cluster_map.groupby(['color', 'color_name']).agg({'position':'count'}).reset_index().rename(columns={"position":"count"}) mydf['Percentage'] = round((mydf['count']/mydf['count'].sum())*100, 1) print(mydf)  # Plotting a scatter plot for all the clusters and their respective colors ax.scatter(cluster_map['x'], cluster_map['y'], cluster_map['z'], color = cluster_map['color']) plt.show()  ''' Subplots with image and a pie chart representing the share of each color identified in the entire photograph/image. ''' plt.figure(figsize=(14, 8)) plt.subplot(221) plt.imshow(img) plt.axis('off')  plt.subplot(222) plt.pie(mydf['count'], labels=mydf['color_name'], colors=mydf['color'], autopct='%1.1f%%', startangle=90) plt.axis('equal') plt.show()
def main(): img = Image.open(IMG_PATH) plotColorClusters(img)