Can now separate image, and write masks to files

py/pallete.py

@@ -0,0 +1,35 @@
+import argparse
+import cv2
+from pallete_schema import PalleteSchema
+
+from os.path import isfile, realpath, basename, dirname, splitext, join
+
+parser = argparse.ArgumentParser(description='Separate an image into most similar colors specified')
+
+parser.add_argument('input', type=str, help='Input image to extract the pallete from')
+
+class Pallete :
+	input = ''
+	output = ''
+	image = None
+
+	def __init__ (self, args) :
+		if isfile(args.input) :
+			self.input = realpath(args.input)
+		else :
+			print(f'File {self.input} does not exist')
+			exit(1)
+		self.set_output()
+
+	def set_output (self) :
+		dir = dirname(self.input)
+		stem = splitext(basename(self.input))[0]
+		self.output = join(dir, f'{stem}.json')
+		print(f'Writing to {stem}.json')
+
+	def process (self) :
+		image = cv2.imread(self.input)
+
+if __name__ == "__main__" :
+	args = parser.parse_args()
+	Pallete(args)

py/pallete_schema.py

@@ -0,0 +1,32 @@
+from jsonschema import validate
+from json import dumps, loads
+
+class PalleteSchema :
+	colors = None
+	schema = {
+		"type" : "array",
+		"items" : {
+			"type" : "object",
+			"properties" : {
+				"name" : { "type" : "string" },
+				"color" : {
+					"type" : "array",
+					"items" : { "type" : "number" }
+				}
+			},
+			"required" : [ "name", "color" ]
+		}
+	}
+	def __init__ (self, file = None):
+		if file is not None:
+			self.parse_file(file)
+
+	def parse_file (self, file) :
+		with open(file) as f :
+			self.parse(f.read())
+		print(f'Parsed pallete file {file}')
+
+	def parse (self, jsonstr) :
+		obj = loads(jsonstr)
+		validate( instance = obj, schema = self.schema)
+		self.colors = obj

py/posterize.py

@@ -6,8 +6,56 @@ import cv2
 class Posterize:
 	"""Posterize an image and then find nearest colors to use"""
 	colors = []
-	#def __init__ (self) :
+	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 + 1
 		
+		self.posterize()
+		self.determine_colors()
+
+	def posterize (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)
+
+		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)
@@ -16,6 +64,3 @@ class Posterize:
 		index_of_smallest = np.where(distances == np.amin(distances))
 		smallest_distance = colors[index_of_smallest]
 		return smallest_distance 
-
-if __name__ == "__main__" :
-	posterize = Posterize()

py/requirements.txt

@@ -1,3 +1,4 @@
 opencv-python
 numpy
-scikit-learn
+scikit-learn
+jsonschema

py/separator.py

@@ -1,18 +1,32 @@
 import argparse
 from posterize import Posterize
-from os.path import abspath
-
-from os.path import isfile
+from pallete_schema import PalleteSchema
+from os.path import isfile, realpath, basename
 
 parser = argparse.ArgumentParser(description='Separate an image into most similar colors specified')
 
 parser.add_argument('input', type=str, help='Input image to separate')
+parser.add_argument('colors', type=int, help='Number of colors to separate into')
+parser.add_argument('pallete', type=str, help='Pallete file')
 parser.add_argument('output', type=str, help='Output dir to write to')
 
 class Separator :
-	posterize = Posterize()
+	input = ''
+	output = ''
+	pallete = None
+
 	def __init__ (self, args) :
-		print(args)
+		if isfile(args.input) :
+			self.input = realpath(args.input)
+		else :
+			print(f'File {args.input} does not exist')
+			exit(1)
+		if isfile(args.pallete) :
+			self.pallete = PalleteSchema(args.pallete)
+		else :
+			print(f'File {args.pallete} does not exist')
+			exit(2)
+		Posterize(self.input, self.pallete, args.colors)
 
 
 if __name__ == "__main__" :

py/test_pallete.json

@@ -0,0 +1,14 @@
+[
+	{ 
+		"name" : "red",
+		"color" : [255, 0, 0]
+	},
+	{ 
+		"name" : "green",
+		"color" : [0, 255, 0]
+	},
+	{ 
+		"name" : "blue",
+		"color" : [0, 0, 255]
+	}
+]