from sklearn.cluster import MiniBatchKMeans import numpy as np import argparse import cv2 class Posterize: """Posterize an image and then find nearest colors to use""" colors = [] image = None pallete = None h = 0 w = 0 n_colors = 3 def __init__ (self, image, pallete, n_colors) : self.image = cv2.imread(image) (self.h, self.w) = self.image.shape[:2] self.pallete = pallete self.n_colors = n_colors self.process() def process (self) : lab = cv2.cvtColor(self.image, cv2.COLOR_BGR2LAB) feature = lab.reshape((self.h * self.w, 3)) clusters = MiniBatchKMeans(n_clusters = self.n_colors, n_init = 'auto') labels = clusters.fit_predict(feature) quant = clusters.cluster_centers_.astype('uint8')[labels] rquant = quant.reshape((self.h, self.w, 3)) rfeature = feature.reshape((self.h, self.w, 3)) bgrquant = cv2.cvtColor(rquant, cv2.COLOR_LAB2BGR) bgrfeature = cv2.cvtColor(rfeature, cv2.COLOR_LAB2BGR) cv2.imshow("image", bgrquant) cv2.waitKey(0) cv2.destroyAllWindows() def closest(self, colors, color): colors = np.array(colors) color = np.array(color) distances = np.sqrt(np.sum((colors - color) ** 2, axis=1)) index_of_smallest = np.where(distances == np.amin(distances)) smallest_distance = colors[index_of_smallest] return smallest_distance