100 lines
2.7 KiB
Python
100 lines
2.7 KiB
Python
from sklearn.cluster import MiniBatchKMeans
|
|
import numpy as np
|
|
import argparse
|
|
import cv2
|
|
from common import convert_color, closest_color, create_colored_image, remove_from_list, list_match
|
|
|
|
class Posterize:
|
|
"""Posterize an image and then find nearest colors to use"""
|
|
colors = []
|
|
colors_dict = {}
|
|
original_colors = []
|
|
|
|
pallete = None
|
|
pallete_space = 'BGR'
|
|
|
|
comparison_space = 'BGR'
|
|
|
|
image = None
|
|
|
|
h = 0
|
|
w = 0
|
|
n_colors = 3
|
|
|
|
white = [255, 255, 255]
|
|
|
|
output = ''
|
|
|
|
def __init__ (self, image, pallete, n_colors, output) :
|
|
self.image = cv2.imread(image)
|
|
(self.h, self.w) = self.image.shape[:2]
|
|
self.pallete = pallete
|
|
self.n_colors = n_colors + 1
|
|
self.output = output
|
|
|
|
self.flatten_pallete()
|
|
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)
|
|
white = closest_color(self.original_colors, [255, 255, 255])
|
|
|
|
blank = create_colored_image(self.w, self.h, [255, 255, 255])
|
|
composite = create_colored_image(self.w, self.h, [255, 255, 255])
|
|
for i in range(self.n_colors) :
|
|
if list_match(self.original_colors[i], white) :
|
|
continue
|
|
mask = self.extract_color_mask(self.image, self.original_colors[i])
|
|
closest = closest_color(self.colors, self.original_colors[i])
|
|
self.colors = remove_from_list(self.colors, closest)
|
|
name = self.match_color_name(closest)
|
|
|
|
cv2.imwrite(f'{name}.png', mask)
|
|
|
|
color_mat = create_colored_image(self.w, self.h, closest)
|
|
color_mask = cv2.bitwise_or(color_mat, color_mat, mask = (255-mask))
|
|
|
|
|
|
|
|
|
|
cv2.imshow('image', composite)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
|
|
def extract_color_mask (self, image, color):
|
|
mask = cv2.inRange(image, color, color)
|
|
return cv2.bitwise_not(mask)
|
|
|
|
def flatten_pallete (self) :
|
|
for color in self.pallete.colors:
|
|
self.colors.append(color['color'])
|
|
self.colors_dict[f'{color["color"][0]},{color["color"][1]},{color["color"][2]}'] = color['name']
|
|
self.pallete_space = color['space']
|
|
|
|
def match_color_name (self, key) :
|
|
return self.colors_dict[f'{key[0]},{key[1]},{key[2]}']
|