HighCWu's picture
init app
801c5bc
import cv2
import numpy as np
import datetime
import pickle
import base64
import json
import gzip
import re
import gradio as gr
from tqdm import tqdm
from cv2.ximgproc import l0Smooth, createGuidedFilter, guidedFilter
import tensorflow
tensorflow.compat.v1.disable_v2_behavior()
tf = tensorflow.compat.v1
import os
import glob
import shutil
splash = glob.glob('ui/web-mobile/splash*')[0]
os.remove(splash)
shutil.copy('res/splash.png', splash)
with open('ui/web-mobile/index.html', 'r', encoding='utf-8') as f:
page = f.read()
with open('ui/web-mobile/index.html', 'w', encoding='utf-8') as f:
f.write(page.replace('Cocos Creator | ', ''))
def ToGray(x):
R = x[:, :, :, 0:1]
G = x[:, :, :, 1:2]
B = x[:, :, :, 2:3]
return 0.30 * R + 0.59 * G + 0.11 * B
def VGG2RGB(x):
return (x + [103.939, 116.779, 123.68])[:, :, :, ::-1]
def norm_feature(x, core):
cs0 = tf.shape(core)[1]
cs1 = tf.shape(core)[2]
small = tf.image.resize_area(x, (cs0, cs1))
avged = tf.nn.avg_pool(tf.pad(small, [[0, 0], [2, 2], [2, 2], [0, 0]], 'REFLECT'), [1, 5, 5, 1], [1, 1, 1, 1],
'VALID')
return tf.image.resize_bicubic(avged, tf.shape(x)[1:3])
def blur(x):
def layer(op):
def layer_decorated(self, *args, **kwargs):
# Automatically set a name if not provided.
name = kwargs.setdefault('name', self.get_unique_name(op.__name__))
# Figure out the layer inputs.
if len(self.terminals) == 0:
raise RuntimeError('No input variables found for layer %s.' % name)
elif len(self.terminals) == 1:
layer_input = self.terminals[0]
else:
layer_input = list(self.terminals)
# Perform the operation and get the output.
layer_output = op(self, layer_input, *args, **kwargs)
# Add to layer LUT.
self.layers[name] = layer_output
# This output is now the input for the next layer.
self.feed(layer_output)
# Return self for chained calls.
return self
return layer_decorated
class Smoother(object):
def __init__(self, inputs, filter_size, sigma):
self.inputs = inputs
self.terminals = []
self.layers = dict(inputs)
self.filter_size = filter_size
self.sigma = sigma
self.setup()
def setup(self):
(self.feed('data')
.conv(name='smoothing'))
def get_unique_name(self, prefix):
ident = sum(t.startswith(prefix) for t, _ in self.layers.items()) + 1
return '%s_%d' % (prefix, ident)
def feed(self, *args):
assert len(args) != 0
self.terminals = []
for fed_layer in args:
if isinstance(fed_layer, str):
try:
fed_layer = self.layers[fed_layer]
except KeyError:
raise KeyError('Unknown layer name fed: %s' % fed_layer)
self.terminals.append(fed_layer)
return self
def gauss_kernel(self, kernlen=21, nsig=3, channels=1):
out_filter = np.load('./nets/gau.npy')
return out_filter
def make_gauss_var(self, name, size, sigma, c_i):
kernel = self.gauss_kernel(size, sigma, c_i)
var = tf.Variable(tf.convert_to_tensor(kernel), name=name)
return var
def get_output(self):
'''Returns the smoother output.'''
return self.terminals[-1]
@layer
def conv(self,
input,
name,
padding='SAME'):
# Get the number of channels in the input
c_i = input.get_shape().as_list()[3]
# Convolution for a given input and kernel
convolve = lambda i, k: tf.nn.depthwise_conv2d(i, k, [1, 1, 1, 1],
padding=padding)
with tf.variable_scope(name) as scope:
kernel = self.make_gauss_var('gauss_weight', self.filter_size,
self.sigma, c_i)
output = convolve(input, kernel)
return output
return Smoother({'data': tf.pad(x, [[0, 0], [9, 9], [9, 9], [0, 0]], 'SYMMETRIC')}, 7, 2).get_output()[:, 9: -9,
9: -9, :]
def downsample(x):
return tf.nn.avg_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
def nts(x):
return (x + [103.939, 116.779, 123.68])[:, :, :, ::-1] / 255.0
def np_expand_image(x):
p = np.pad(x, ((1, 1), (1, 1), (0, 0)), 'symmetric')
r = []
r.append(p[:-2, 1:-1, :])
r.append(p[1:-1, :-2, :])
r.append(p[1:-1, 1:-1, :])
r.append(p[1:-1, 2:, :])
r.append(p[2:, 1:-1, :])
return np.stack(r, axis=2)
def build_sketch_sparse(x, abs):
x = x[:, :, None].astype(np.float32)
expanded = np_expand_image(x)
distance = x[:, :, None] - expanded
if abs:
distance = np.abs(distance)
weight = 8 - distance
weight[weight < 0] = 0.0
weight /= np.sum(weight, axis=2, keepdims=True)
return weight
def build_repeat_mulsep(x, m, i):
a = m[:, :, 0]
b = m[:, :, 1]
c = m[:, :, 2]
d = m[:, :, 3]
e = m[:, :, 4]
y = x
for _ in range(i):
p = tf.pad(y, [[1, 1], [1, 1], [0, 0]], 'SYMMETRIC')
y = p[:-2, 1:-1, :] * a + p[1:-1, :-2, :] * b + y * c + p[1:-1, 2:, :] * d + p[2:, 1:-1, :] * e
return y
session = tf.Session()
tf.keras.backend.set_session(session)
ip1 = tf.placeholder(dtype=tf.float32, shape=(None, None, None, 1))
ip3 = tf.placeholder(dtype=tf.float32, shape=(None, None, None, 3))
ip4 = tf.placeholder(dtype=tf.float32, shape=(None, None, None, 4))
ipsp9 = tf.placeholder(dtype=tf.float32, shape=(None, None, 5, 1))
ipsp3 = tf.placeholder(dtype=tf.float32, shape=(None, None, 3))
tf_sparse_op_H = build_repeat_mulsep(ipsp3, ipsp9, 64)
tf_sparse_op_L = build_repeat_mulsep(ipsp3, ipsp9, 16)
def make_graph():
with gzip.open('./nets/refs.net', 'rb') as fp:
refs_img = pickle.load(fp)
tail = tf.keras.models.load_model('./nets/tail.net')
reader = tf.keras.models.load_model('./nets/reader.net')
head = tf.keras.models.load_model('./nets/head.net')
neck = tf.keras.models.load_model('./nets/neck.net')
inception = tf.keras.models.load_model('./nets/inception.net')
render_head = tf.keras.models.load_model('./nets/render_head.net')
render_neck = tf.keras.models.load_model('./nets/render_neck.net')
tail_op = tail(ip3)
features = reader(ip3 / 255.0)
print('Loaded some basic models.')
feed = [1 - ip1 / 255.0, (ip4[:, :, :, 0:3] / 127.5 - 1) * ip4[:, :, :, 3:4] / 255.0]
for _ in range(len(features)):
feed.append(tf.reduce_mean(features[_], axis=[1, 2]))
nil0, nil1, head_temp = head(feed)
feed[0] = tf.clip_by_value(1 - tf.image.resize_bilinear(ToGray(VGG2RGB(head_temp) / 255.0), tf.shape(ip1)[1:3]),
0.0, 1.0)
nil4, nil5, head_temp = neck(feed)
head_op = VGG2RGB(head_temp)
features_render = inception((ip3 + (downsample(ip1) - blur(downsample(ip1))) * 2.0) / 255.0)
precessed_feed = [(ip4[:, :, :, 0:3] / 127.5 - 1) * ip4[:, :, :, 3:4] / 255.0] + [
norm_feature(item, features_render[-1]) for item in features_render]
nil6, nil7, render_A = render_head([1 - ip1 / 255.0] + precessed_feed)
nil8, nil9, render_B = render_neck(
[1 - tf.image.resize_bilinear(ToGray(nts(render_A)), tf.shape(ip1)[1:3])] + precessed_feed)
render_op = nts(render_B) * 255.0
print('Loaded - Style2Paints Deep Learning Engine V4.6 - GPU')
session.run(tf.global_variables_initializer())
tail.load_weights('./nets/tail.net')
head.load_weights('./nets/head.net')
neck.load_weights('./nets/neck.net')
reader.load_weights('./nets/reader.net')
inception.load_weights('./nets/inception.net')
render_head.load_weights('./nets/render_head.net')
render_neck.load_weights('./nets/render_neck.net')
print('Deep learning modules are ready.')
return tail_op, head_op, render_op, refs_img
tail_op_g, head_op_g, render_op_g, refs_img_g = make_graph()
def go_tail(x):
def srange(l, s):
result = []
iters = int(float(l) / float(s))
for i in range(iters):
result.append([i * s, (i + 1) * s])
result[len(result) - 1][1] = l
return result
H, W, C = x.shape
padded_img = np.pad(x, ((20, 20), (20, 20), (0, 0)), 'symmetric').astype(np.float32) / 255.0
lines = []
for hs, he in srange(H, 64):
items = []
for ws, we in srange(W, 64):
items.append(padded_img[hs:he + 40, ws:we + 40, :])
lines.append(items)
iex = 0
result_all_lines = []
for line in lines:
result_one_line = []
for item in line:
ots = session.run(tail_op_g, feed_dict={ip3: item[None, :, :, :]})[0]
result_one_line.append(ots[41:-41, 41:-41, :])
print('Slicing ... ' + str(iex))
iex += 1
result_one_line = np.concatenate(result_one_line, axis=1)
result_all_lines.append(result_one_line)
result_all_lines = np.concatenate(result_all_lines, axis=0)
return (result_all_lines * 255.0).clip(0, 255).astype(np.uint8)
def go_head(sketch, global_hint, local_hint):
return session.run(head_op_g, feed_dict={
ip1: sketch[None, :, :, None], ip3: global_hint[None, :, :, :], ip4: local_hint[None, :, :, :]
})[0].clip(0, 255).astype(np.uint8)
def go_render(sketch, segmentation, points):
return session.run(render_op_g, feed_dict={
ip1: sketch[None, :, :, None], ip3: segmentation[None, :, :, :], ip4: points[None, :, :, :]
})[0].clip(0, 255).astype(np.uint8)
print('Deep learning functions are ready.')
def k_resize(x, k):
if x.shape[0] < x.shape[1]:
s0 = k
s1 = int(x.shape[1] * (k / x.shape[0]))
s1 = s1 - s1 % 64
_s0 = 16 * s0
_s1 = int(x.shape[1] * (_s0 / x.shape[0]))
_s1 = (_s1 + 32) - (_s1 + 32) % 64
else:
s1 = k
s0 = int(x.shape[0] * (k / x.shape[1]))
s0 = s0 - s0 % 64
_s1 = 16 * s1
_s0 = int(x.shape[0] * (_s1 / x.shape[1]))
_s0 = (_s0 + 32) - (_s0 + 32) % 64
new_min = min(_s1, _s0)
raw_min = min(x.shape[0], x.shape[1])
if new_min < raw_min:
interpolation = cv2.INTER_AREA
else:
interpolation = cv2.INTER_LANCZOS4
y = cv2.resize(x, (_s1, _s0), interpolation=interpolation)
return y
def d_resize(x, d, fac=1.0):
new_min = min(int(d[1] * fac), int(d[0] * fac))
raw_min = min(x.shape[0], x.shape[1])
if new_min < raw_min:
interpolation = cv2.INTER_AREA
else:
interpolation = cv2.INTER_LANCZOS4
y = cv2.resize(x, (int(d[1] * fac), int(d[0] * fac)), interpolation=interpolation)
return y
def min_resize(x, m):
if x.shape[0] < x.shape[1]:
s0 = m
s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
else:
s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
s1 = m
new_max = min(s1, s0)
raw_max = min(x.shape[0], x.shape[1])
if new_max < raw_max:
interpolation = cv2.INTER_AREA
else:
interpolation = cv2.INTER_LANCZOS4
y = cv2.resize(x, (s1, s0), interpolation=interpolation)
return y
def cli_norm(sketch):
light = np.max(min_resize(sketch, 64), axis=(0, 1), keepdims=True)
intensity = (light - sketch.astype(np.float32)).clip(0, 255)
line_intensities = np.sort(intensity[intensity > 16])[::-1]
line_quantity = float(line_intensities.shape[0])
intensity /= line_intensities[int(line_quantity * 0.1)]
intensity *= 0.9
return (255.0 - intensity * 255.0).clip(0, 255).astype(np.uint8)
def cv2_imwrite(a, b):
print(a)
cv2.imwrite(a, b)
def from_png_to_jpg(map):
if map.shape[2] == 3:
return map
color = map[:, :, 0:3].astype(np.float) / 255.0
alpha = map[:, :, 3:4].astype(np.float) / 255.0
reversed_color = 1 - color
final_color = (255.0 - reversed_color * alpha * 255.0).clip(0, 255).astype(np.uint8)
return final_color
def s_enhance(x, k=2.0):
p = cv2.cvtColor(x, cv2.COLOR_RGB2HSV).astype(np.float)
p[:, :, 1] *= k
p = p.clip(0, 255).astype(np.uint8)
return cv2.cvtColor(p, cv2.COLOR_HSV2RGB).clip(0, 255)
def ini_hint(x):
r = np.zeros(shape=(x.shape[0], x.shape[1], 4), dtype=np.uint8)
return r
def opreate_normal_hint(gird, points, length):
h = gird.shape[0]
w = gird.shape[1]
for point in points:
x, y, r, g, b = point
x = int(x * w)
y = int(y * h)
l_ = max(0, x - length)
b_ = max(0, y - length)
r_ = min(w, x + length + 1)
t_ = min(h, y + length + 1)
gird[b_:t_, l_:r_, 2] = r
gird[b_:t_, l_:r_, 1] = g
gird[b_:t_, l_:r_, 0] = b
gird[b_:t_, l_:r_, 3] = 255.0
return gird
def get_hdr(x):
def get_hdr_g(x):
img = x.astype(np.float32)
mean = np.mean(img)
h_mean = mean.copy()
l_mean = mean.copy()
for i in range(2):
h_mean = np.mean(img[img >= h_mean])
l_mean = np.mean(img[img <= l_mean])
for i in range(2):
l_mean = np.mean(img[img <= l_mean])
return l_mean, mean, h_mean
l_mean = np.zeros(shape=(1, 1, 3), dtype=np.float32)
mean = np.zeros(shape=(1, 1, 3), dtype=np.float32)
h_mean = np.zeros(shape=(1, 1, 3), dtype=np.float32)
for c in range(3):
l, m, h = get_hdr_g(x[:, :, c])
l_mean[:, :, c] = l
mean[:, :, c] = m
h_mean[:, :, c] = h
return l_mean, mean, h_mean
def f2(x1, x2, x3, y1, y2, y3, x):
A = y1 * ((x - x2) * (x - x3)) / ((x1 - x2) * (x1 - x3))
B = y2 * ((x - x1) * (x - x3)) / ((x2 - x1) * (x2 - x3))
C = y3 * ((x - x1) * (x - x2)) / ((x3 - x1) * (x3 - x2))
return A + B + C
print('Tricks loaded.')
def refine_image(image, sketch, origin):
verbose = False
def cv_log(name, img):
if verbose:
print(name)
cv2.imshow('cv_log', img.clip(0, 255).astype(np.uint8))
cv2.imwrite('cv_log.png', img.clip(0, 255).astype(np.uint8))
cv2.waitKey(0)
print('Building Sparse Matrix ...')
sketch = sketch.astype(np.float32)
sparse_matrix = build_sketch_sparse(sketch, True)
bright_matrix = build_sketch_sparse(sketch - cv2.GaussianBlur(sketch, (0, 0), 3.0), False)
guided_matrix = createGuidedFilter(sketch.clip(0, 255).astype(np.uint8), 1, 0.01)
HDRL, HDRM, HDRH = get_hdr(image)
def go_guide(x):
y = x + (x - cv2.GaussianBlur(x, (0, 0), 1)) * 2.0
for _ in tqdm(range(4)):
y = guided_matrix.filter(y)
return y
def go_refine_sparse(x):
return session.run(tf_sparse_op_H, feed_dict={ipsp3: x, ipsp9: sparse_matrix})
def go_refine_bright(x):
return session.run(tf_sparse_op_L, feed_dict={ipsp3: x, ipsp9: bright_matrix})
def go_flat(x):
pia = 32
y = x.clip(0, 255).astype(np.uint8)
y = cv2.resize(y, (x.shape[1] // 2, x.shape[0] // 2), interpolation=cv2.INTER_AREA)
y = np.pad(y, ((pia, pia), (pia, pia), (0, 0)), 'reflect')
y = l0Smooth(y, None, 0.01)
y = y[pia:-pia, pia:-pia, :]
y = cv2.resize(y, (x.shape[1], x.shape[0]), interpolation=cv2.INTER_CUBIC)
return y
def go_hdr(x):
xl, xm, xh = get_hdr(x)
y = f2(xl, xm, xh, HDRL, HDRM, HDRH, x)
return y.clip(0, 255)
def go_blend(BGR, X, m):
BGR = BGR.clip(0, 255).astype(np.uint8)
X = X.clip(0, 255).astype(np.uint8)
YUV = cv2.cvtColor(BGR, cv2.COLOR_BGR2YUV)
s_l = YUV[:, :, 0].astype(np.float32)
t_l = X.astype(np.float32)
r_l = (s_l * t_l / 255.0) if m else np.minimum(s_l, t_l)
YUV[:, :, 0] = r_l.clip(0, 255).astype(np.uint8)
return cv2.cvtColor(YUV, cv2.COLOR_YUV2BGR)
print('Getting Target ...')
smoothed = d_resize(image, sketch.shape)
print('Global Optimization ...')
cv_log('smoothed', smoothed)
sparse_smoothed = go_refine_sparse(smoothed)
cv_log('smoothed', sparse_smoothed)
smoothed = go_guide(sparse_smoothed)
cv_log('smoothed', smoothed)
smoothed = go_hdr(smoothed)
cv_log('smoothed', smoothed)
print('Decomposition Optimization ...')
flat = sparse_smoothed.copy()
cv_log('flat', flat)
flat = go_refine_bright(flat)
cv_log('flat', flat)
flat = go_flat(flat)
cv_log('flat', flat)
flat = go_refine_sparse(flat)
cv_log('flat', flat)
flat = go_guide(flat)
cv_log('flat', flat)
flat = go_hdr(flat)
cv_log('flat', flat)
print('Blending Optimization ...')
cv_log('origin', origin)
blended_smoothed = go_blend(smoothed, origin, False)
cv_log('blended_smoothed', blended_smoothed)
blended_flat = go_blend(flat, origin, True)
cv_log('blended_flat', blended_flat)
print('Optimization finished.')
return smoothed, flat, blended_smoothed, blended_flat
print('Fundamental Methods loaded.')
def cv2_encode(image: np.ndarray):
if image is None:
return 'null'
_, data = cv2.imencode('.png', image)
return 'data:image/png;base64,' + base64.b64encode(data).decode('utf8')
def get_request_image(request, name):
img = request.get(name)
img = re.sub('^data:image/.+;base64,', '', img)
img = base64.b64decode(img)
img = np.fromstring(img, dtype=np.uint8)
img = cv2.imdecode(img, -1)
return img
def upload_sketch(json_str, history):
request = json.loads(json_str)
origin = from_png_to_jpg(get_request_image(request, 'sketch'))
ID = datetime.datetime.now().strftime('H%HM%MS%S')
print('New room ID: ' + ID)
sketch = min_resize(origin, 512)
sketch = np.min(sketch, axis=2)
sketch = cli_norm(sketch)
sketch = np.tile(sketch[:, :, None], [1, 1, 3])
sketch = go_tail(sketch)
sketch = np.mean(sketch, axis=2)
if len(history) > 5:
history = history[-5:]
return ID + '_' + ID, [*history, {'ID': ID, 'origin': origin, 'sketch': sketch, 'results': {}}]
def request_result(json_str, history):
request = json.loads(json_str)
room = request.get("room")
ridx = next(i for i, his in enumerate(history) if his['ID'] == room)
room_path = './rooms/' + room + '/'
ID = datetime.datetime.now().strftime('H%HM%MS%S')
history[ridx]['results'][ID] = {}
points = request.get("points")
points = json.loads(points)
history[ridx]['results'][ID]['points'] = points
for _ in range(len(points)):
points[_][1] = 1 - points[_][1]
sketch = history[ridx]['sketch']
origin = history[ridx]['origin']
if origin.ndim == 3:
origin = cv2.cvtColor(origin, cv2.COLOR_BGR2GRAY)
origin = d_resize(origin, sketch.shape).astype(np.float32)
low_origin = cv2.GaussianBlur(origin, (0, 0), 3.0)
high_origin = origin - low_origin
low_origin = (low_origin / np.median(low_origin) * 255.0).clip(0, 255)
origin = (low_origin + high_origin).clip(0, 255).astype(np.uint8)
faceID = int(request.get("faceID")) - 65535
print(faceID)
if faceID > -1:
print('Default reference.')
face = from_png_to_jpg(refs_img_g[faceID])
else:
print('Load reference.')
face = from_png_to_jpg(get_request_image(request, 'face'))
face = s_enhance(face, 2.0)
print('request result room = ' + str(room) + ', ID = ' + str(ID))
print('processing painting in ' + room_path)
sketch_1024 = k_resize(sketch, 64)
hints_1024 = opreate_normal_hint(ini_hint(sketch_1024), points, length=2)
careless = go_head(sketch_1024, k_resize(face, 14), hints_1024)
smoothed_careless, flat_careless, blended_smoothed_careless, blended_flat_careless = refine_image(careless, sketch,
origin)
history[ridx]['results'][ID]['smoothed_careless'] = smoothed_careless
history[ridx]['results'][ID]['flat_careless'] = flat_careless
history[ridx]['results'][ID]['blended_smoothed_careless'] = blended_smoothed_careless
history[ridx]['results'][ID]['blended_flat_careless'] = blended_flat_careless
print('Stage I finished.')
careful = go_render(sketch_1024, d_resize(flat_careless, sketch_1024.shape, 0.5), hints_1024)
smoothed_careful, flat_careful, blended_smoothed_careful, blended_flat_careful = refine_image(careful, sketch,
origin)
history[ridx]['results'][ID]['smoothed_careful'] = smoothed_careful
history[ridx]['results'][ID]['flat_careful'] = flat_careful
history[ridx]['results'][ID]['blended_smoothed_careful'] = blended_smoothed_careful
history[ridx]['results'][ID]['blended_flat_careful'] = blended_flat_careful
history[ridx]['results'][ID]['lighted'] = blended_flat_careful
print('Stage II finished.')
return room + '_' + ID, history
def download_result(json_str, history):
request = json.loads(json_str)
room = request.get("room")
step = request.get("step")
name = request.get("name")
ridx = next(i for i, his in enumerate(history) if his['ID'] == room)
if history[ridx].get(name, None) is None:
result = history[ridx]['results'][step][name]
else:
result = history[ridx][name]
if name == 'points':
return json.dumps(result)
return cv2_encode(result)
with gr.Blocks() as demo:
history = gr.State(value=[])
with gr.Row():
with gr.Column():
btn_show = gr.Button("Open Style2Paints V4.6")
btn_show.click(None, _js="(_) => open('file/ui/web-mobile/index.html')")
with gr.Row():
with gr.Box():
with gr.Row():
upload_sketch_json = gr.Textbox(label="upload_sketch(json string)")
with gr.Row():
upload_sketch_btn = gr.Button(label="Submit sketch json")
with gr.Row():
upload_sketch_result = gr.Textbox(label="Result", interactive=False)
upload_sketch_btn.click(upload_sketch, [upload_sketch_json, history], [upload_sketch_result, history], api_name="upload_sketch")
with gr.Box():
with gr.Row():
request_result_json = gr.Textbox(label="request_result(json string)")
with gr.Row():
request_result_btn = gr.Button(label="Submit json of request for result")
with gr.Row():
request_result_result = gr.Textbox(label="Result", interactive=False)
upload_sketch_btn.click(request_result, [request_result_json, history], [request_result_result, history], api_name="request_result")
with gr.Box():
with gr.Row():
download_result_json = gr.Textbox(label="download_result(json string)")
with gr.Row():
download_result_btn = gr.Button(label="Submit json of download for result")
with gr.Row():
download_result_result = gr.Textbox(label="Result", interactive=False)
upload_sketch_btn.click(download_result, [download_result_json, history], [download_result_result], api_name="download_result")
demo.launch()