From f862e28a8bf645ee3cd49edff1b3f46fe2e1f4cf Mon Sep 17 00:00:00 2001
From: mattmcw <matt@sixteenmillimeter.com>
Date: Sat, 21 Oct 2023 13:51:38 -0400
Subject: [PATCH] Now posterizes an image to a number of colors determined by
 the command line arguments on separator.py. Still needs to color match and
 then export to different layer files.

---
 py/pallete.py        | 35 +++++++++++++++++++++++++++++++++++
 py/pallete_schema.py | 32 ++++++++++++++++++++++++++++++++
 py/posterize.py      | 35 +++++++++++++++++++++++++++++++----
 py/requirements.txt  |  3 ++-
 py/separator.py      | 24 +++++++++++++++++++-----
 py/test_pallete.json | 14 ++++++++++++++
 6 files changed, 133 insertions(+), 10 deletions(-)
 create mode 100644 py/pallete.py
 create mode 100644 py/pallete_schema.py
 create mode 100644 py/test_pallete.json

diff --git a/py/pallete.py b/py/pallete.py
new file mode 100644
index 0000000..6f4f290
--- /dev/null
+++ b/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)
diff --git a/py/pallete_schema.py b/py/pallete_schema.py
new file mode 100644
index 0000000..d6af7ec
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/py/posterize.py b/py/posterize.py
index c8168c5..00f6ed2 100644
--- a/py/posterize.py
+++ b/py/posterize.py
@@ -6,8 +6,38 @@ import cv2
 class Posterize:
 	"""Posterize an image and then find nearest colors to use"""
 	colors = []
-	#def __init__ (self) :
+	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)
@@ -16,6 +46,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()
\ No newline at end of file
diff --git a/py/requirements.txt b/py/requirements.txt
index 6e95bc6..ec3e673 100644
--- a/py/requirements.txt
+++ b/py/requirements.txt
@@ -1,3 +1,4 @@
 opencv-python
 numpy
-scikit-learn
\ No newline at end of file
+scikit-learn
+jsonschema
\ No newline at end of file
diff --git a/py/separator.py b/py/separator.py
index e0481dd..cb09be7 100644
--- a/py/separator.py
+++ b/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__" :
diff --git a/py/test_pallete.json b/py/test_pallete.json
new file mode 100644
index 0000000..a887c64
--- /dev/null
+++ b/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]
+	}
+]
\ No newline at end of file