268 lines
7.6 KiB
Python
268 lines
7.6 KiB
Python
from sklearn.cluster import MiniBatchKMeans
|
|
import numpy as np
|
|
import argparse
|
|
import cv2
|
|
from common import convert_color, closest_color_weighted_euclidean, closest_color_euclidean, create_colored_image, remove_from_list, list_match, to_luma, convertScale
|
|
import os
|
|
import subprocess
|
|
from multiprocessing.pool import ThreadPool
|
|
from functools import cmp_to_key
|
|
|
|
def sort_pallete (a, b):
|
|
A = to_luma(a, 'BGR')
|
|
B = to_luma(b, 'BGR')
|
|
if A > B :
|
|
return 1
|
|
if B < A :
|
|
return -1
|
|
return 0
|
|
|
|
class Posterize:
|
|
"""Posterize an image and then find nearest colors to use"""
|
|
colors = []
|
|
colors_dict = {}
|
|
original_colors = []
|
|
layers = []
|
|
previews = []
|
|
svgs = []
|
|
headless = False
|
|
jobs = 1
|
|
|
|
pallete = None
|
|
pallete_space = 'BGR'
|
|
comparison_space = 'BGR'
|
|
|
|
image = None
|
|
|
|
h = 0
|
|
w = 0
|
|
n_colors = 3
|
|
max_particles = 17000
|
|
conf = os.path.abspath('./conf/base.conf')
|
|
stipple_gen = os.path.abspath('../../stipple_gen')
|
|
|
|
white = [255, 255, 255]
|
|
|
|
output = None
|
|
|
|
def __init__ (self, image, pallete, n_colors, output, headless, jobs, particles) :
|
|
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.headless = headless
|
|
self.jobs = jobs
|
|
self.max_particles = particles
|
|
|
|
if not os.path.exists(self.output) :
|
|
print(f'Output directory {self.output} does not exist, creating...')
|
|
os.makedirs(self.output)
|
|
|
|
self.flatten_pallete()
|
|
self.posterize()
|
|
self.determine_colors()
|
|
self.stipple()
|
|
self.preview()
|
|
self.optimize()
|
|
|
|
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
|
|
|
|
self.show(bgrquant)
|
|
|
|
def determine_colors (self):
|
|
reshaped = self.image.reshape(-1, self.image.shape[2])
|
|
self.original_colors = np.unique(reshaped, axis=0)
|
|
white, white_dist = closest_color_weighted_euclidean(self.original_colors, [255, 255, 255], 'BGR')
|
|
|
|
blank = create_colored_image(self.w, self.h, [255, 255, 255])
|
|
composite = create_colored_image(self.w, self.h, [255, 255, 255])
|
|
|
|
mask = self.extract_color_mask(self.image, white)
|
|
layer_name = f'WHITE.png'
|
|
output_layer = os.path.join(self.output, layer_name)
|
|
cv2.imwrite(output_layer, mask)
|
|
|
|
self.layers.append({
|
|
'layer' : output_layer,
|
|
'color' : white,
|
|
'space' : self.pallete_space
|
|
})
|
|
|
|
for i in range(self.n_colors) :
|
|
if list_match(self.original_colors[i], white) :
|
|
continue
|
|
original = self.original_colors[i] #BGR
|
|
mask = self.extract_color_mask(self.image, original)
|
|
original_normalized = convert_color(original, 'BGR', self.pallete_space)
|
|
if self.pallete_space == 'RGB' or self.pallete_space == 'BGR' :
|
|
closest, dist = closest_color_weighted_euclidean(self.colors, original_normalized, self.pallete_space)
|
|
else :
|
|
closest, dist = closest_color_euclidean(self.colors, original_normalized)
|
|
self.colors = remove_from_list(self.colors, closest)
|
|
name = self.match_color_name(closest)
|
|
layer_name = f'{name}.png'
|
|
output_layer = os.path.join(self.output, layer_name)
|
|
cv2.imwrite(output_layer, mask)
|
|
self.layers.append({
|
|
'layer' : output_layer,
|
|
'color' : closest,
|
|
'space' : self.pallete_space
|
|
})
|
|
mask = cv2.bitwise_not(mask)
|
|
composite[mask > 0] = np.array(closest)
|
|
|
|
composite_name = f'posterized.png'
|
|
composite_path = os.path.join(self.output, composite_name)
|
|
|
|
cv2.imwrite(composite_path, composite)
|
|
|
|
self.show(composite)
|
|
|
|
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']
|
|
|
|
#self.colors = sorted(self.colors, key=cmp_to_key(sort_pallete))
|
|
#for color in self.colors :
|
|
# print(to_luma(color, self.pallete_space))
|
|
|
|
#quit()
|
|
|
|
def match_color_name (self, key) :
|
|
return self.colors_dict[f'{key[0]},{key[1]},{key[2]}']
|
|
|
|
def stipple (self) :
|
|
sanity_check = 0
|
|
cmds = []
|
|
|
|
for layer in self.layers :
|
|
if 'WHITE.png' in layer['layer'] :
|
|
continue
|
|
l = cv2.imread(layer['layer'], 0)
|
|
(h, w) = l.shape[:2]
|
|
total = h * w
|
|
black = total - cv2.countNonZero(l)
|
|
ratio = black/total
|
|
max_particles = round(ratio * self.max_particles)
|
|
input_image = os.path.abspath(layer['layer'])
|
|
dir_name, file_name = os.path.split(input_image)
|
|
file_part, ext = os.path.splitext(file_name)
|
|
output_image = os.path.join(dir_name, f'{file_part}_preview.png')
|
|
output_svg = os.path.join(dir_name, f'{file_part}.svg')
|
|
cmd = [
|
|
'bash',
|
|
'stipple_gen.sh',
|
|
'--inputImage', input_image,
|
|
'--outputImage', output_image,
|
|
'--outputSVG', output_svg,
|
|
'--config', self.conf,
|
|
'--maxParticles', str(max_particles),
|
|
'--canvasWidth', str(self.w),
|
|
'--canvasHeight', str(self.h),
|
|
'--windowWidth', str(self.w),
|
|
'--windowHeight', str(self.h)
|
|
]
|
|
cmds.append({
|
|
'cmd' : cmd,
|
|
'color' : layer['color'],
|
|
'output_image' : output_image,
|
|
'output_svg' : output_svg
|
|
})
|
|
|
|
if self.jobs > 1 and len(cmds) > 1:
|
|
print(f'Running {self.jobs} stipple_gen processes simultaneously for {len(cmds)} jobs')
|
|
results = []
|
|
pool = ThreadPool(self.jobs)
|
|
for job in cmds :
|
|
result = pool.apply_async(self.render, (job,))
|
|
results.append(result)
|
|
[result.wait() for result in results]
|
|
else :
|
|
for job in cmds :
|
|
self.render(job)
|
|
|
|
|
|
def render (self, job) :
|
|
name = os.path.basename(job['output_svg'])
|
|
print(f'Starting stipple_gen {name}...')
|
|
proc = subprocess.Popen(job['cmd'], cwd = self.stipple_gen, stdout=subprocess.PIPE)
|
|
while True:
|
|
line = proc.stdout.readline()
|
|
if not line:
|
|
break
|
|
l = line.decode('ascii').strip()
|
|
print(f'[{name}] {l}')
|
|
self.svgs.append(job['output_svg'])
|
|
self.previews.append({
|
|
'layer' : job['output_image'],
|
|
'color' : job['color']
|
|
})
|
|
|
|
def preview (self) :
|
|
composite = create_colored_image(self.w, self.h, [255, 255, 255])
|
|
for layer in self.previews :
|
|
print(f'Compositing {layer["layer"]}')
|
|
l = cv2.imread(layer['layer'], 0)
|
|
mask = cv2.bitwise_not(l)
|
|
composite[mask > 0] = np.array(convert_color(layer['color'], self.pallete_space, 'BGR'))
|
|
|
|
composite_name = f'preview.png'
|
|
composite_path = os.path.join(self.output, composite_name)
|
|
|
|
cv2.imwrite(composite_path, composite)
|
|
|
|
self.show(composite)
|
|
|
|
def optimize (self) :
|
|
if self.jobs > 1 and len(self.svgs) > 1:
|
|
print(f'Running {self.jobs} svgopt processes simultaneously for {len(self.svgs)} jobs')
|
|
results = []
|
|
pool = ThreadPool(self.jobs)
|
|
for svg in self.svgs:
|
|
result = pool.apply_async(self.svgopt, (svg,))
|
|
results.append(result)
|
|
[result.wait() for result in results]
|
|
else :
|
|
for svg in self.svgs :
|
|
self.svgopt(svg)
|
|
|
|
def svgopt (self, svg) :
|
|
name = os.path.basename(svg)
|
|
cmd = [ 'svgopt', svg, svg]
|
|
print(f'Optimizing {name}...')
|
|
proc = subprocess.Popen(cmd, cwd = self.stipple_gen, stdout=subprocess.PIPE)
|
|
while True:
|
|
line = proc.stdout.readline()
|
|
if not line:
|
|
break
|
|
l = line.decode('ascii').strip()
|
|
print(f'[{name}] {l}')
|
|
|
|
def show (self, mat) :
|
|
if not self.headless :
|
|
cv2.imshow('image', mat)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|