Can now separate image, and write masks to files

py/posterize.py

@@ -6,20 +6,24 @@ import cv2
 class Posterize:
 	"""Posterize an image and then find nearest colors to use"""
 	colors = []
+	original_colors = []
 	image = None
 	pallete = None
 	h = 0
 	w = 0
 	n_colors = 3
+	white = [255, 255, 255]
 
 	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.n_colors = n_colors + 1
-		self.process()
 		
-	def process (self) :
+		self.posterize()
+		self.determine_colors()
+
+	def posterize (self):
 		lab = cv2.cvtColor(self.image, cv2.COLOR_BGR2LAB)
 		feature = lab.reshape((self.h * self.w, 3))
 		
@@ -32,12 +36,26 @@ class Posterize:
 		rfeature = feature.reshape((self.h, self.w, 3))
 
 		bgrquant = cv2.cvtColor(rquant, cv2.COLOR_LAB2BGR)
-		bgrfeature = cv2.cvtColor(rfeature, cv2.COLOR_LAB2BGR)
+		#bgrfeature = cv2.cvtColor(rfeature, cv2.COLOR_LAB2BGR)
+
+		self.image = bgrquant
+		
 
 		cv2.imshow("image", bgrquant)
 		cv2.waitKey(0)
 		cv2.destroyAllWindows()
 
+	def determine_colors (self):
+		reshaped = self.image.reshape(-1, self.image.shape[2])
+		self.original_colors = np.unique(reshaped, axis=0)
+		print(self.original_colors)
+		for i in range(self.n_colors) :
+			mask = self.extract_color(self.image, self.original_colors[i])
+			cv2.imwrite(f'{i}.png', mask)
+
+	def extract_color (self, image, color):
+		mask = cv2.inRange(image, color, color)
+		return cv2.bitwise_not(mask)
 
 	def closest(self, colors, color):
 		colors = np.array(colors)