marker_separation/py/posterize.py

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