Compare commits

..

2 Commits

6 changed files with 151 additions and 10 deletions

35
py/pallete.py Normal file
View File

@ -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)

32
py/pallete_schema.py Normal file
View File

@ -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

View File

@ -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()

View File

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

View File

@ -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__" :

14
py/test_pallete.json Normal file
View File

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