49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
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
|