Spaces:
Runtime error
Runtime error
jaekookang
commited on
Commit
•
9ff1108
1
Parent(s):
a871eb8
first upload
Browse files- .gitignore +4 -0
- README.md +2 -1
- brush/brush_large_horizontal.png +0 -0
- brush/brush_large_vertical.png +0 -0
- gradio_painttransformer.py +86 -0
- inference.py +72 -0
- input/abst1.jpg +0 -0
- input/cat.jpg +0 -0
- input/kanagawa.jpg +0 -0
- input/obama.jpg +0 -0
- input/van_gogh_starry_night.jpg +0 -0
- network.py +74 -0
- render_parallel.py +245 -0
- render_serial.py +283 -0
- render_utils.py +102 -0
- requirements.txt +8 -0
.gitignore
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
__pycache__
|
2 |
+
*.log
|
3 |
+
*.pdparams
|
4 |
+
*.mp4
|
README.md
CHANGED
@@ -4,7 +4,7 @@ emoji: 🏃
|
|
4 |
colorFrom: purple
|
5 |
colorTo: gray
|
6 |
sdk: gradio
|
7 |
-
app_file:
|
8 |
pinned: false
|
9 |
---
|
10 |
|
@@ -35,3 +35,4 @@ Path is relative to the root of the repository.
|
|
35 |
|
36 |
`pinned`: _boolean_
|
37 |
Whether the Space stays on top of your list.
|
|
|
|
4 |
colorFrom: purple
|
5 |
colorTo: gray
|
6 |
sdk: gradio
|
7 |
+
app_file: gradio_painttransformer.py
|
8 |
pinned: false
|
9 |
---
|
10 |
|
|
|
35 |
|
36 |
`pinned`: _boolean_
|
37 |
Whether the Space stays on top of your list.
|
38 |
+
|
brush/brush_large_horizontal.png
ADDED
brush/brush_large_vertical.png
ADDED
gradio_painttransformer.py
ADDED
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''PaintTransformer Demo
|
2 |
+
|
3 |
+
- 2021-12-21 first created
|
4 |
+
- See: https://github.com/wzmsltw/PaintTransformer
|
5 |
+
|
6 |
+
'''
|
7 |
+
|
8 |
+
import os
|
9 |
+
import cv2
|
10 |
+
import network
|
11 |
+
from time import time
|
12 |
+
from glob import glob
|
13 |
+
from loguru import logger
|
14 |
+
import gradio as gr
|
15 |
+
|
16 |
+
import paddle
|
17 |
+
import render_utils
|
18 |
+
import render_parallel
|
19 |
+
import render_serial
|
20 |
+
|
21 |
+
# ---------- Settings ----------
|
22 |
+
GPU_ID = '-1'
|
23 |
+
os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID
|
24 |
+
DEVICE = 'cpu' if GPU_ID == '-1' else f'cuda:{GPU_ID}'
|
25 |
+
|
26 |
+
examples = sorted(glob(os.path.join('input', '*.jpg')))
|
27 |
+
WIDTH = 512
|
28 |
+
HEIGHT = 512
|
29 |
+
STROKE_NUM = 8
|
30 |
+
FPS = 10
|
31 |
+
|
32 |
+
# ---------- Logger ----------
|
33 |
+
logger.add('app.log', mode='a')
|
34 |
+
logger.info('===== APP RESTARTED =====')
|
35 |
+
|
36 |
+
# ---------- Model ----------
|
37 |
+
MODEL_FILE = 'paint_best.pdparams'
|
38 |
+
if not os.path.exists(MODEL_FILE):
|
39 |
+
os.system('gdown --id 1G0O81qSvGp0kFCgyaQHmPygbVHFi1--q')
|
40 |
+
logger.info('model downloaded')
|
41 |
+
else:
|
42 |
+
logger.info('model already exists')
|
43 |
+
|
44 |
+
paddle.set_device(DEVICE)
|
45 |
+
net_g = network.Painter(5, STROKE_NUM, 256, 8, 3, 3)
|
46 |
+
net_g.set_state_dict(paddle.load(MODEL_FILE))
|
47 |
+
net_g.eval()
|
48 |
+
for param in net_g.parameters():
|
49 |
+
param.stop_gradient = True
|
50 |
+
|
51 |
+
brush_large_vertical = render_utils.read_img('brush/brush_large_vertical.png', 'L')
|
52 |
+
brush_large_horizontal = render_utils.read_img('brush/brush_large_horizontal.png', 'L')
|
53 |
+
meta_brushes = paddle.concat([brush_large_vertical, brush_large_horizontal], axis=0)
|
54 |
+
|
55 |
+
def predict(image_file):
|
56 |
+
original_img = render_utils.read_img(image_file, 'RGB', WIDTH, HEIGHT)
|
57 |
+
logger.info(f'--- image loaded & resized {WIDTH}x{HEIGHT}')
|
58 |
+
|
59 |
+
logger.info('--- doing inference...')
|
60 |
+
t0 = time()
|
61 |
+
final_result_list = render_serial.render_serial(original_img, net_g, meta_brushes)
|
62 |
+
logger.info(f'--- inference took {time() - t0:.4f} sec')
|
63 |
+
|
64 |
+
out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), FPS,
|
65 |
+
(WIDTH, HEIGHT))
|
66 |
+
for idx, frame in enumerate(final_result_list):
|
67 |
+
out.write(frame)
|
68 |
+
out.release()
|
69 |
+
logger.info('--- animation generated')
|
70 |
+
return 'output.mp4'
|
71 |
+
|
72 |
+
iface = gr.Interface(
|
73 |
+
predict,
|
74 |
+
title='🎨 Paint Transformer',
|
75 |
+
description='This demo converts an image into a sequence of painted images (animation)',
|
76 |
+
inputs=[
|
77 |
+
gr.inputs.Image(label='Input image', type='filepath')
|
78 |
+
],
|
79 |
+
outputs=[
|
80 |
+
gr.outputs.Video(label='Output animation', type='mp4')
|
81 |
+
],
|
82 |
+
examples=examples,
|
83 |
+
article='<p style="text-align:center">Original work: <a href="https://github.com/wzmsltw/PaintTransformer">PaintTransformer</a></p>'
|
84 |
+
)
|
85 |
+
|
86 |
+
iface.launch(debug=True)
|
inference.py
ADDED
@@ -0,0 +1,72 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn.functional as F
|
3 |
+
import numpy as np
|
4 |
+
from PIL import Image
|
5 |
+
import network
|
6 |
+
import os
|
7 |
+
import math
|
8 |
+
import render_utils
|
9 |
+
import paddle
|
10 |
+
import paddle.nn as nn
|
11 |
+
import paddle.nn.functional as F
|
12 |
+
import cv2
|
13 |
+
import render_parallel
|
14 |
+
import render_serial
|
15 |
+
|
16 |
+
def main(input_path, model_path, output_dir, need_animation=False, resize_h=None, resize_w=None, serial=False):
|
17 |
+
if not os.path.exists(output_dir):
|
18 |
+
os.mkdir(output_dir)
|
19 |
+
input_name = os.path.basename(input_path)
|
20 |
+
output_path = os.path.join(output_dir, input_name)
|
21 |
+
frame_dir = None
|
22 |
+
if need_animation:
|
23 |
+
if not serial:
|
24 |
+
print('It must be under serial mode if animation results are required, so serial flag is set to True!')
|
25 |
+
serial = True
|
26 |
+
frame_dir = os.path.join(output_dir, input_name[:input_name.find('.')])
|
27 |
+
if not os.path.exists(frame_dir):
|
28 |
+
os.mkdir(frame_dir)
|
29 |
+
stroke_num = 8
|
30 |
+
|
31 |
+
#* ----- load model ----- *#
|
32 |
+
# paddle.set_device('gpu')
|
33 |
+
paddle.set_device('cpu') # 2021-12-21 jkang edited to "cpu"
|
34 |
+
net_g = network.Painter(5, stroke_num, 256, 8, 3, 3)
|
35 |
+
net_g.set_state_dict(paddle.load(model_path))
|
36 |
+
net_g.eval()
|
37 |
+
for param in net_g.parameters():
|
38 |
+
param.stop_gradient = True
|
39 |
+
|
40 |
+
#* ----- load brush ----- *#
|
41 |
+
brush_large_vertical = render_utils.read_img('brush/brush_large_vertical.png', 'L')
|
42 |
+
brush_large_horizontal = render_utils.read_img('brush/brush_large_horizontal.png', 'L')
|
43 |
+
meta_brushes = paddle.concat([brush_large_vertical, brush_large_horizontal], axis=0)
|
44 |
+
|
45 |
+
import time
|
46 |
+
t0 = time.time()
|
47 |
+
|
48 |
+
original_img = render_utils.read_img(input_path, 'RGB', resize_h, resize_w)
|
49 |
+
if serial:
|
50 |
+
final_result_list = render_serial.render_serial(original_img, net_g, meta_brushes)
|
51 |
+
if need_animation:
|
52 |
+
|
53 |
+
print("total frame:", len(final_result_list))
|
54 |
+
for idx, frame in enumerate(final_result_list):
|
55 |
+
cv2.imwrite(os.path.join(frame_dir, '%03d.png' %idx), frame)
|
56 |
+
else:
|
57 |
+
cv2.imwrite(output_path, final_result_list[-1])
|
58 |
+
else:
|
59 |
+
final_result = render_parallel.render_parallel(original_img, net_g, meta_brushes)
|
60 |
+
cv2.imwrite(output_path, final_result)
|
61 |
+
|
62 |
+
print("total infer time:", time.time() - t0)
|
63 |
+
|
64 |
+
if __name__ == '__main__':
|
65 |
+
|
66 |
+
main(input_path='input/chicago.jpg',
|
67 |
+
model_path='paint_best.pdparams',
|
68 |
+
output_dir='output/',
|
69 |
+
need_animation=True, # whether need intermediate results for animation.
|
70 |
+
resize_h=512, # resize original input to this size. None means do not resize.
|
71 |
+
resize_w=512, # resize original input to this size. None means do not resize.
|
72 |
+
serial=True) # if need animation, serial must be True.
|
input/abst1.jpg
ADDED
input/cat.jpg
ADDED
input/kanagawa.jpg
ADDED
input/obama.jpg
ADDED
input/van_gogh_starry_night.jpg
ADDED
network.py
ADDED
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import paddle
|
2 |
+
import paddle.nn as nn
|
3 |
+
import math
|
4 |
+
|
5 |
+
class Painter(nn.Layer):
|
6 |
+
"""
|
7 |
+
network architecture written in paddle.
|
8 |
+
"""
|
9 |
+
def __init__(self, param_per_stroke, total_strokes, hidden_dim, n_heads=8, n_enc_layers=3, n_dec_layers=3):
|
10 |
+
super().__init__()
|
11 |
+
self.enc_img = nn.Sequential(
|
12 |
+
nn.Pad2D([1, 1, 1, 1], 'reflect'),
|
13 |
+
nn.Conv2D(3, 32, 3, 1),
|
14 |
+
nn.BatchNorm2D(32),
|
15 |
+
nn.ReLU(), # maybe replace with the inplace version
|
16 |
+
nn.Pad2D([1, 1, 1, 1], 'reflect'),
|
17 |
+
nn.Conv2D(32, 64, 3, 2),
|
18 |
+
nn.BatchNorm2D(64),
|
19 |
+
nn.ReLU(),
|
20 |
+
nn.Pad2D([1, 1, 1, 1], 'reflect'),
|
21 |
+
nn.Conv2D(64, 128, 3, 2),
|
22 |
+
nn.BatchNorm2D(128),
|
23 |
+
nn.ReLU())
|
24 |
+
self.enc_canvas = nn.Sequential(
|
25 |
+
nn.Pad2D([1, 1, 1, 1], 'reflect'),
|
26 |
+
nn.Conv2D(3, 32, 3, 1),
|
27 |
+
nn.BatchNorm2D(32),
|
28 |
+
nn.ReLU(),
|
29 |
+
nn.Pad2D([1, 1, 1, 1], 'reflect'),
|
30 |
+
nn.Conv2D(32, 64, 3, 2),
|
31 |
+
nn.BatchNorm2D(64),
|
32 |
+
nn.ReLU(),
|
33 |
+
nn.Pad2D([1, 1, 1, 1], 'reflect'),
|
34 |
+
nn.Conv2D(64, 128, 3, 2),
|
35 |
+
nn.BatchNorm2D(128),
|
36 |
+
nn.ReLU())
|
37 |
+
self.conv = nn.Conv2D(128 * 2, hidden_dim, 1)
|
38 |
+
self.transformer = nn.Transformer(hidden_dim, n_heads, n_enc_layers, n_dec_layers)
|
39 |
+
self.linear_param = nn.Sequential(
|
40 |
+
nn.Linear(hidden_dim, hidden_dim),
|
41 |
+
nn.ReLU(),
|
42 |
+
nn.Linear(hidden_dim, hidden_dim),
|
43 |
+
nn.ReLU(),
|
44 |
+
nn.Linear(hidden_dim, param_per_stroke))
|
45 |
+
self.linear_decider = nn.Linear(hidden_dim, 1)
|
46 |
+
self.query_pos = paddle.static.create_parameter([total_strokes, hidden_dim], dtype='float32',
|
47 |
+
default_initializer=nn.initializer.Uniform(0, 1))
|
48 |
+
self.row_embed = paddle.static.create_parameter([8, hidden_dim // 2], dtype='float32',
|
49 |
+
default_initializer=nn.initializer.Uniform(0, 1))
|
50 |
+
self.col_embed = paddle.static.create_parameter([8, hidden_dim // 2], dtype='float32',
|
51 |
+
default_initializer=nn.initializer.Uniform(0, 1))
|
52 |
+
|
53 |
+
def forward(self, img, canvas):
|
54 |
+
"""
|
55 |
+
prediction
|
56 |
+
"""
|
57 |
+
b, _, H, W = img.shape
|
58 |
+
img_feat = self.enc_img(img)
|
59 |
+
canvas_feat = self.enc_canvas(canvas)
|
60 |
+
h, w = img_feat.shape[-2:]
|
61 |
+
feat = paddle.concat([img_feat, canvas_feat], axis=1)
|
62 |
+
feat_conv = self.conv(feat)
|
63 |
+
|
64 |
+
pos_embed = paddle.concat([
|
65 |
+
self.col_embed[:w].unsqueeze(0).tile([h, 1, 1]),
|
66 |
+
self.row_embed[:h].unsqueeze(1).tile([1, w, 1]),
|
67 |
+
], axis=-1).flatten(0, 1).unsqueeze(1)
|
68 |
+
|
69 |
+
hidden_state = self.transformer((pos_embed + feat_conv.flatten(2).transpose([2, 0, 1])).transpose([1, 0, 2]),
|
70 |
+
self.query_pos.unsqueeze(1).tile([1, b, 1]).transpose([1, 0, 2]))
|
71 |
+
|
72 |
+
param = self.linear_param(hidden_state)
|
73 |
+
decision = self.linear_decider(hidden_state)
|
74 |
+
return param, decision
|
render_parallel.py
ADDED
@@ -0,0 +1,245 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import render_utils
|
2 |
+
import paddle
|
3 |
+
import paddle.nn as nn
|
4 |
+
import paddle.nn.functional as F
|
5 |
+
import numpy as np
|
6 |
+
import math
|
7 |
+
|
8 |
+
def crop(img, h, w):
|
9 |
+
H, W = img.shape[-2:]
|
10 |
+
pad_h = (H - h) // 2
|
11 |
+
pad_w = (W - w) // 2
|
12 |
+
remainder_h = (H - h) % 2
|
13 |
+
remainder_w = (W - w) % 2
|
14 |
+
img = img[:, :, pad_h:H - pad_h - remainder_h, pad_w:W - pad_w - remainder_w]
|
15 |
+
return img
|
16 |
+
|
17 |
+
def stroke_net_predict(img_patch, result_patch, patch_size, net_g, stroke_num, patch_num):
|
18 |
+
"""
|
19 |
+
stroke_net_predict
|
20 |
+
"""
|
21 |
+
img_patch = img_patch.transpose([0, 2, 1]).reshape([-1, 3, patch_size, patch_size])
|
22 |
+
result_patch = result_patch.transpose([0, 2, 1]).reshape([-1, 3, patch_size, patch_size])
|
23 |
+
#*----- Stroke Predictor -----*#
|
24 |
+
shape_param, stroke_decision = net_g(img_patch, result_patch)
|
25 |
+
stroke_decision = (stroke_decision > 0).astype('float32')
|
26 |
+
#*----- sampling color -----*#
|
27 |
+
grid = shape_param[:, :, :2].reshape([img_patch.shape[0] * stroke_num, 1, 1, 2])
|
28 |
+
img_temp = img_patch.unsqueeze(1).tile([1, stroke_num, 1, 1, 1]).reshape([
|
29 |
+
img_patch.shape[0] * stroke_num, 3, patch_size, patch_size])
|
30 |
+
color = nn.functional.grid_sample(img_temp, 2 * grid - 1, align_corners=False).reshape([
|
31 |
+
img_patch.shape[0], stroke_num, 3])
|
32 |
+
param = paddle.concat([shape_param, color], axis=-1)
|
33 |
+
|
34 |
+
param = param.reshape([-1, 8])
|
35 |
+
param[:, :2] = param[:, :2] / 2 + 0.25
|
36 |
+
param[:, 2:4] = param[:, 2:4] / 2
|
37 |
+
param = param.reshape([1, patch_num, patch_num, stroke_num, 8])
|
38 |
+
decision = stroke_decision.reshape([1, patch_num, patch_num, stroke_num])#.astype('bool')
|
39 |
+
return param, decision
|
40 |
+
|
41 |
+
|
42 |
+
def param2img_parallel(param, decision, meta_brushes, cur_canvas, stroke_num=8):
|
43 |
+
"""
|
44 |
+
Input stroke parameters and decisions for each patch, meta brushes, current canvas, frame directory,
|
45 |
+
and whether there is a border (if intermediate painting results are required).
|
46 |
+
Output the painting results of adding the corresponding strokes on the current canvas.
|
47 |
+
Args:
|
48 |
+
param: a tensor with shape batch size x patch along height dimension x patch along width dimension
|
49 |
+
x n_stroke_per_patch x n_param_per_stroke
|
50 |
+
decision: a 01 tensor with shape batch size x patch along height dimension x patch along width dimension
|
51 |
+
x n_stroke_per_patch
|
52 |
+
meta_brushes: a tensor with shape 2 x 3 x meta_brush_height x meta_brush_width.
|
53 |
+
The first slice on the batch dimension denotes vertical brush and the second one denotes horizontal brush.
|
54 |
+
cur_canvas: a tensor with shape batch size x 3 x H x W,
|
55 |
+
where H and W denote height and width of padded results of original images.
|
56 |
+
|
57 |
+
Returns:
|
58 |
+
cur_canvas: a tensor with shape batch size x 3 x H x W, denoting painting results.
|
59 |
+
"""
|
60 |
+
# param: b, h, w, stroke_per_patch, param_per_stroke
|
61 |
+
# decision: b, h, w, stroke_per_patch
|
62 |
+
b, h, w, s, p = param.shape
|
63 |
+
h, w = int(h), int(w)
|
64 |
+
param = param.reshape([-1, 8])
|
65 |
+
decision = decision.reshape([-1, 8])
|
66 |
+
|
67 |
+
H, W = cur_canvas.shape[-2:]
|
68 |
+
is_odd_y = h % 2 == 1
|
69 |
+
is_odd_x = w % 2 == 1
|
70 |
+
render_size_y = 2 * H // h
|
71 |
+
render_size_x = 2 * W // w
|
72 |
+
|
73 |
+
even_idx_y = paddle.arange(0, h, 2)
|
74 |
+
even_idx_x = paddle.arange(0, w, 2)
|
75 |
+
if h > 1:
|
76 |
+
odd_idx_y = paddle.arange(1, h, 2)
|
77 |
+
if w > 1:
|
78 |
+
odd_idx_x = paddle.arange(1, w, 2)
|
79 |
+
|
80 |
+
cur_canvas = F.pad(cur_canvas, [render_size_x // 4, render_size_x // 4,
|
81 |
+
render_size_y // 4, render_size_y // 4])
|
82 |
+
|
83 |
+
valid_foregrounds = render_utils.param2stroke(param, render_size_y, render_size_x, meta_brushes)
|
84 |
+
|
85 |
+
#* ----- load dilation/erosion ---- *#
|
86 |
+
dilation = render_utils.Dilation2d(m=1)
|
87 |
+
erosion = render_utils.Erosion2d(m=1)
|
88 |
+
|
89 |
+
#* ----- generate alphas ----- *#
|
90 |
+
valid_alphas = (valid_foregrounds > 0).astype('float32')
|
91 |
+
valid_foregrounds = valid_foregrounds.reshape([-1, stroke_num, 1, render_size_y, render_size_x])
|
92 |
+
valid_alphas = valid_alphas.reshape([-1, stroke_num, 1, render_size_y, render_size_x])
|
93 |
+
|
94 |
+
temp = [dilation(valid_foregrounds[:, i, :, :, :]) for i in range(stroke_num)]
|
95 |
+
valid_foregrounds = paddle.stack(temp, axis=1)
|
96 |
+
valid_foregrounds = valid_foregrounds.reshape([-1, 1, render_size_y, render_size_x])
|
97 |
+
|
98 |
+
temp = [erosion(valid_alphas[:, i, :, :, :]) for i in range(stroke_num)]
|
99 |
+
valid_alphas = paddle.stack(temp, axis=1)
|
100 |
+
valid_alphas = valid_alphas.reshape([-1, 1, render_size_y, render_size_x])
|
101 |
+
|
102 |
+
foregrounds = valid_foregrounds.reshape([-1, h, w, stroke_num, 1, render_size_y, render_size_x])
|
103 |
+
alphas = valid_alphas.reshape([-1, h, w, stroke_num, 1, render_size_y, render_size_x])
|
104 |
+
decision = decision.reshape([-1, h, w, stroke_num, 1, 1, 1])
|
105 |
+
param = param.reshape([-1, h, w, stroke_num, 8])
|
106 |
+
|
107 |
+
def partial_render(this_canvas, patch_coord_y, patch_coord_x):
|
108 |
+
canvas_patch = F.unfold(this_canvas, [render_size_y, render_size_x], strides=[render_size_y // 2, render_size_x // 2])
|
109 |
+
# canvas_patch: b, 3 * py * px, h * w
|
110 |
+
canvas_patch = canvas_patch.reshape([b, 3, render_size_y, render_size_x, h, w])
|
111 |
+
canvas_patch = canvas_patch.transpose([0, 4, 5, 1, 2, 3])
|
112 |
+
selected_canvas_patch = paddle.gather(canvas_patch, patch_coord_y, 1)
|
113 |
+
selected_canvas_patch = paddle.gather(selected_canvas_patch, patch_coord_x, 2)
|
114 |
+
selected_canvas_patch = selected_canvas_patch.reshape([0, 0, 0, 1, 3, render_size_y, render_size_x])
|
115 |
+
selected_foregrounds = paddle.gather(foregrounds, patch_coord_y, 1)
|
116 |
+
selected_foregrounds = paddle.gather(selected_foregrounds, patch_coord_x, 2)
|
117 |
+
selected_alphas = paddle.gather(alphas, patch_coord_y, 1)
|
118 |
+
selected_alphas = paddle.gather(selected_alphas, patch_coord_x, 2)
|
119 |
+
selected_decisions = paddle.gather(decision, patch_coord_y, 1)
|
120 |
+
selected_decisions = paddle.gather(selected_decisions, patch_coord_x, 2)
|
121 |
+
selected_color = paddle.gather(param, patch_coord_y, 1)
|
122 |
+
selected_color = paddle.gather(selected_color, patch_coord_x, 2)
|
123 |
+
selected_color = paddle.gather(selected_color, paddle.to_tensor([5,6,7]), 4)
|
124 |
+
selected_color = selected_color.reshape([0, 0, 0, stroke_num, 3, 1, 1])
|
125 |
+
|
126 |
+
for i in range(stroke_num):
|
127 |
+
i = paddle.to_tensor(i)
|
128 |
+
|
129 |
+
cur_foreground = paddle.gather(selected_foregrounds, i, 3)
|
130 |
+
cur_alpha = paddle.gather(selected_alphas, i, 3)
|
131 |
+
cur_decision = paddle.gather(selected_decisions, i, 3)
|
132 |
+
cur_color = paddle.gather(selected_color, i, 3)
|
133 |
+
cur_foreground = cur_foreground * cur_color
|
134 |
+
selected_canvas_patch = cur_foreground * cur_alpha * cur_decision + selected_canvas_patch * (1 - cur_alpha * cur_decision)
|
135 |
+
|
136 |
+
selected_canvas_patch = selected_canvas_patch.reshape([0, 0, 0, 3, render_size_y, render_size_x])
|
137 |
+
this_canvas = selected_canvas_patch.transpose([0, 3, 1, 4, 2, 5])
|
138 |
+
|
139 |
+
# this_canvas: b, 3, h_half, py, w_half, px
|
140 |
+
h_half = this_canvas.shape[2]
|
141 |
+
w_half = this_canvas.shape[4]
|
142 |
+
this_canvas = this_canvas.reshape([b, 3, h_half * render_size_y, w_half * render_size_x])
|
143 |
+
# this_canvas: b, 3, h_half * py, w_half * px
|
144 |
+
return this_canvas
|
145 |
+
|
146 |
+
# even - even area
|
147 |
+
# 1 | 0
|
148 |
+
# 0 | 0
|
149 |
+
canvas = partial_render(cur_canvas, even_idx_y, even_idx_x)
|
150 |
+
if not is_odd_y:
|
151 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, -render_size_y // 2:, :canvas.shape[3]]], axis=2)
|
152 |
+
if not is_odd_x:
|
153 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, :canvas.shape[2], -render_size_x // 2:]], axis=3)
|
154 |
+
cur_canvas = canvas
|
155 |
+
|
156 |
+
# odd - odd area
|
157 |
+
# 0 | 0
|
158 |
+
# 0 | 1
|
159 |
+
if h > 1 and w > 1:
|
160 |
+
canvas = partial_render(cur_canvas, odd_idx_y, odd_idx_x)
|
161 |
+
canvas = paddle.concat([cur_canvas[:, :, :render_size_y // 2, -canvas.shape[3]:], canvas], axis=2)
|
162 |
+
canvas = paddle.concat([cur_canvas[:, :, -canvas.shape[2]:, :render_size_x // 2], canvas], axis=3)
|
163 |
+
if is_odd_y:
|
164 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, -render_size_y // 2:, :canvas.shape[3]]], axis=2)
|
165 |
+
if is_odd_x:
|
166 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, :canvas.shape[2], -render_size_x // 2:]], axis=3)
|
167 |
+
cur_canvas = canvas
|
168 |
+
|
169 |
+
# odd - even area
|
170 |
+
# 0 | 0
|
171 |
+
# 1 | 0
|
172 |
+
if h > 1:
|
173 |
+
canvas = partial_render(cur_canvas, odd_idx_y, even_idx_x)
|
174 |
+
canvas = paddle.concat([cur_canvas[:, :, :render_size_y // 2, :canvas.shape[3]], canvas], axis=2)
|
175 |
+
if is_odd_y:
|
176 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, -render_size_y // 2:, :canvas.shape[3]]], axis=2)
|
177 |
+
if not is_odd_x:
|
178 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, :canvas.shape[2], -render_size_x // 2:]], axis=3)
|
179 |
+
cur_canvas = canvas
|
180 |
+
|
181 |
+
# odd - even area
|
182 |
+
# 0 | 1
|
183 |
+
# 0 | 0
|
184 |
+
if w > 1:
|
185 |
+
canvas = partial_render(cur_canvas, even_idx_y, odd_idx_x)
|
186 |
+
canvas = paddle.concat([cur_canvas[:, :, :canvas.shape[2], :render_size_x // 2], canvas], axis=3)
|
187 |
+
if not is_odd_y:
|
188 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, -render_size_y // 2:, -canvas.shape[3]:]], axis=2)
|
189 |
+
if is_odd_x:
|
190 |
+
canvas = paddle.concat([canvas, cur_canvas[:, :, :canvas.shape[2], -render_size_x // 2:]], axis=3)
|
191 |
+
cur_canvas = canvas
|
192 |
+
|
193 |
+
cur_canvas = cur_canvas[:, :, render_size_y // 4:-render_size_y // 4, render_size_x // 4:-render_size_x // 4]
|
194 |
+
|
195 |
+
return cur_canvas
|
196 |
+
|
197 |
+
|
198 |
+
|
199 |
+
def render_parallel(original_img, net_g, meta_brushes):
|
200 |
+
|
201 |
+
patch_size = 32
|
202 |
+
stroke_num = 8
|
203 |
+
|
204 |
+
with paddle.no_grad():
|
205 |
+
|
206 |
+
original_h, original_w = original_img.shape[-2:]
|
207 |
+
K = max(math.ceil(math.log2(max(original_h, original_w) / patch_size)), 0)
|
208 |
+
original_img_pad_size = patch_size * (2 ** K)
|
209 |
+
original_img_pad = render_utils.pad(original_img, original_img_pad_size, original_img_pad_size)
|
210 |
+
final_result = paddle.zeros_like(original_img)
|
211 |
+
|
212 |
+
for layer in range(0, K + 1):
|
213 |
+
layer_size = patch_size * (2 ** layer)
|
214 |
+
|
215 |
+
img = F.interpolate(original_img_pad, (layer_size, layer_size))
|
216 |
+
result = F.interpolate(final_result, (layer_size, layer_size))
|
217 |
+
img_patch = F.unfold(img, [patch_size, patch_size], strides=[patch_size, patch_size])
|
218 |
+
result_patch = F.unfold(result, [patch_size, patch_size], strides=[patch_size, patch_size])
|
219 |
+
|
220 |
+
# There are patch_num * patch_num patches in total
|
221 |
+
patch_num = (layer_size - patch_size) // patch_size + 1
|
222 |
+
param, decision = stroke_net_predict(img_patch, result_patch, patch_size, net_g, stroke_num, patch_num)
|
223 |
+
|
224 |
+
#print(param.shape, decision.shape)
|
225 |
+
final_result = param2img_parallel(param, decision, meta_brushes, final_result)
|
226 |
+
|
227 |
+
# paint another time for last layer
|
228 |
+
border_size = original_img_pad_size // (2 * patch_num)
|
229 |
+
img = F.interpolate(original_img_pad, (layer_size, layer_size))
|
230 |
+
result = F.interpolate(final_result, (layer_size, layer_size))
|
231 |
+
img = F.pad(img, [patch_size // 2, patch_size // 2, patch_size // 2, patch_size // 2])
|
232 |
+
result = F.pad(result, [patch_size // 2, patch_size // 2, patch_size // 2, patch_size // 2])
|
233 |
+
img_patch = F.unfold(img, [patch_size, patch_size], strides=[patch_size, patch_size])
|
234 |
+
result_patch = F.unfold(result, [patch_size, patch_size], strides=[patch_size, patch_size])
|
235 |
+
final_result = F.pad(final_result, [border_size, border_size, border_size, border_size])
|
236 |
+
patch_num = (img.shape[2] - patch_size) // patch_size + 1
|
237 |
+
#w = (img.shape[3] - patch_size) // patch_size + 1
|
238 |
+
|
239 |
+
param, decision = stroke_net_predict(img_patch, result_patch, patch_size, net_g, stroke_num, patch_num)
|
240 |
+
|
241 |
+
final_result = param2img_parallel(param, decision, meta_brushes, final_result)
|
242 |
+
|
243 |
+
final_result = final_result[:, :, border_size:-border_size, border_size:-border_size]
|
244 |
+
final_result = (final_result.numpy().squeeze().transpose([1,2,0])[:,:,::-1] * 255).astype(np.uint8)
|
245 |
+
return final_result
|
render_serial.py
ADDED
@@ -0,0 +1,283 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# !/usr/bin/env python3
|
2 |
+
"""
|
3 |
+
codes for oilpainting style transfer.
|
4 |
+
"""
|
5 |
+
import paddle
|
6 |
+
import paddle.nn as nn
|
7 |
+
import paddle.nn.functional as F
|
8 |
+
import numpy as np
|
9 |
+
from PIL import Image
|
10 |
+
import math
|
11 |
+
import cv2
|
12 |
+
import render_utils
|
13 |
+
import time
|
14 |
+
|
15 |
+
|
16 |
+
def get_single_layer_lists(param, decision, ori_img, render_size_x, render_size_y, h, w, meta_brushes, dilation, erosion, stroke_num):
|
17 |
+
"""
|
18 |
+
get_single_layer_lists
|
19 |
+
"""
|
20 |
+
valid_foregrounds = render_utils.param2stroke(param[:, :], render_size_y, render_size_x, meta_brushes)
|
21 |
+
|
22 |
+
valid_alphas = (valid_foregrounds > 0).astype('float32')
|
23 |
+
valid_foregrounds = valid_foregrounds.reshape([-1, stroke_num, 1, render_size_y, render_size_x])
|
24 |
+
valid_alphas = valid_alphas.reshape([-1, stroke_num, 1, render_size_y, render_size_x])
|
25 |
+
|
26 |
+
temp = [dilation(valid_foregrounds[:, i, :, :, :]) for i in range(stroke_num)]
|
27 |
+
valid_foregrounds = paddle.stack(temp, axis=1)
|
28 |
+
valid_foregrounds = valid_foregrounds.reshape([-1, 1, render_size_y, render_size_x])
|
29 |
+
|
30 |
+
temp = [erosion(valid_alphas[:, i, :, :, :]) for i in range(stroke_num)]
|
31 |
+
valid_alphas = paddle.stack(temp, axis=1)
|
32 |
+
valid_alphas = valid_alphas.reshape([-1, 1, render_size_y, render_size_x])
|
33 |
+
|
34 |
+
patch_y = 4 * render_size_y // 5
|
35 |
+
patch_x = 4 * render_size_x // 5
|
36 |
+
|
37 |
+
img_patch = ori_img.reshape([1, 3, h, ori_img.shape[2]//h, w, ori_img.shape[3]//w])
|
38 |
+
img_patch = img_patch.transpose([0, 2, 4, 1, 3, 5])[0]
|
39 |
+
|
40 |
+
xid_list = []
|
41 |
+
yid_list = []
|
42 |
+
error_list = []
|
43 |
+
|
44 |
+
for flag_idx, flag in enumerate(decision.cpu().numpy()):
|
45 |
+
if flag:
|
46 |
+
flag_idx = flag_idx // stroke_num
|
47 |
+
x_id = flag_idx % w
|
48 |
+
flag_idx = flag_idx // w
|
49 |
+
y_id = flag_idx % h
|
50 |
+
xid_list.append(x_id)
|
51 |
+
yid_list.append(y_id)
|
52 |
+
|
53 |
+
inner_fores = valid_foregrounds[:, :, render_size_y // 10:9 * render_size_y // 10,
|
54 |
+
render_size_x // 10:9 * render_size_x // 10]
|
55 |
+
inner_alpha = valid_alphas[:, :, render_size_y // 10:9 * render_size_y // 10,
|
56 |
+
render_size_x // 10:9 * render_size_x // 10]
|
57 |
+
inner_fores = inner_fores.reshape([h * w, stroke_num, 1, patch_y, patch_x])
|
58 |
+
inner_alpha = inner_alpha.reshape([h * w, stroke_num, 1, patch_y, patch_x])
|
59 |
+
inner_real = img_patch.reshape([h * w, 3, patch_y, patch_x]).unsqueeze(1)
|
60 |
+
|
61 |
+
R = param[:, 5]
|
62 |
+
G = param[:, 6]
|
63 |
+
B = param[:, 7]#, G, B = param[5:]
|
64 |
+
R = R.reshape([-1, stroke_num]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
|
65 |
+
G = G.reshape([-1, stroke_num]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
|
66 |
+
B = B.reshape([-1, stroke_num]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
|
67 |
+
error_R = R * inner_fores - inner_real[:, :, 0:1, :, :]
|
68 |
+
error_G = G * inner_fores - inner_real[:, :, 1:2, :, :]
|
69 |
+
error_B = B * inner_fores - inner_real[:, :, 2:3, :, :]
|
70 |
+
error = paddle.abs(error_R) + paddle.abs(error_G)+ paddle.abs(error_B)
|
71 |
+
|
72 |
+
error = error * inner_alpha
|
73 |
+
error = paddle.sum(error, axis=(2, 3, 4)) / paddle.sum(inner_alpha, axis=(2, 3, 4))
|
74 |
+
error_list = error.reshape([-1]).numpy()[decision.numpy()]
|
75 |
+
error_list = list(error_list)
|
76 |
+
|
77 |
+
valid_foregrounds = paddle.to_tensor(valid_foregrounds.numpy()[decision.numpy()])
|
78 |
+
valid_alphas = paddle.to_tensor(valid_alphas.numpy()[decision.numpy()])
|
79 |
+
|
80 |
+
selected_param = paddle.to_tensor(param.numpy()[decision.numpy()])
|
81 |
+
return xid_list, yid_list, valid_foregrounds, valid_alphas, error_list, selected_param
|
82 |
+
|
83 |
+
|
84 |
+
def get_single_stroke_on_full_image_A(x_id, y_id, valid_foregrounds, valid_alphas, param, original_img,
|
85 |
+
render_size_x, render_size_y, patch_x, patch_y):
|
86 |
+
"""
|
87 |
+
get_single_stroke_on_full_image_A
|
88 |
+
"""
|
89 |
+
tmp_foreground = paddle.zeros_like(original_img)
|
90 |
+
|
91 |
+
patch_y_num = original_img.shape[2] // patch_y
|
92 |
+
patch_x_num = original_img.shape[3] // patch_x
|
93 |
+
|
94 |
+
brush = valid_foregrounds.unsqueeze(0)
|
95 |
+
color_map = param[5:]
|
96 |
+
brush = brush.tile([1, 3, 1, 1])
|
97 |
+
color_map = color_map.unsqueeze(-1).unsqueeze(-1).unsqueeze(0)#.repeat(1, 1, H, W)
|
98 |
+
brush = brush * color_map
|
99 |
+
|
100 |
+
pad_l = x_id * patch_x
|
101 |
+
pad_r = (patch_x_num - x_id - 1) * patch_x
|
102 |
+
pad_t = y_id * patch_y
|
103 |
+
pad_b = (patch_y_num - y_id - 1) * patch_y
|
104 |
+
tmp_foreground = nn.functional.pad(brush, [pad_l, pad_r, pad_t, pad_b])
|
105 |
+
tmp_foreground = tmp_foreground[:, :, render_size_y // 10:-render_size_y // 10,
|
106 |
+
render_size_x // 10:-render_size_x // 10]
|
107 |
+
|
108 |
+
tmp_alpha = nn.functional.pad(valid_alphas.unsqueeze(0), [pad_l, pad_r, pad_t, pad_b])
|
109 |
+
tmp_alpha = tmp_alpha[:, :, render_size_y // 10:-render_size_y // 10, render_size_x // 10:-render_size_x // 10]
|
110 |
+
return tmp_foreground, tmp_alpha
|
111 |
+
|
112 |
+
def get_single_stroke_on_full_image_B(x_id, y_id, valid_foregrounds, valid_alphas, param,
|
113 |
+
original_img, render_size_x, render_size_y, patch_x, patch_y):
|
114 |
+
"""
|
115 |
+
get_single_stroke_on_full_image_B
|
116 |
+
"""
|
117 |
+
x_expand = patch_x // 2 + render_size_x // 10
|
118 |
+
y_expand = patch_y // 2 + render_size_y // 10
|
119 |
+
|
120 |
+
pad_l = x_id * patch_x
|
121 |
+
pad_r = original_img.shape[3] + 2 * x_expand - (x_id * patch_x + render_size_x)
|
122 |
+
pad_t = y_id * patch_y
|
123 |
+
pad_b = original_img.shape[2] + 2 * y_expand - (y_id * patch_y + render_size_y)
|
124 |
+
|
125 |
+
brush = valid_foregrounds.unsqueeze(0)
|
126 |
+
color_map = param[5:]
|
127 |
+
brush = brush.tile([1, 3, 1, 1])
|
128 |
+
color_map = color_map.unsqueeze(-1).unsqueeze(-1).unsqueeze(0)#.repeat(1, 1, H, W)
|
129 |
+
brush = brush * color_map
|
130 |
+
|
131 |
+
tmp_foreground = nn.functional.pad(brush, [pad_l, pad_r, pad_t, pad_b])
|
132 |
+
|
133 |
+
tmp_foreground = tmp_foreground[:, :, y_expand:- y_expand, x_expand:-x_expand]
|
134 |
+
tmp_alpha = nn.functional.pad(valid_alphas.unsqueeze(0), [pad_l, pad_r, pad_t, pad_b])
|
135 |
+
tmp_alpha = tmp_alpha[:, :, y_expand:- y_expand, x_expand:-x_expand]
|
136 |
+
return tmp_foreground, tmp_alpha
|
137 |
+
|
138 |
+
def stroke_net_predict(img_patch, result_patch, patch_size, net_g, stroke_num):
|
139 |
+
"""
|
140 |
+
stroke_net_predict
|
141 |
+
"""
|
142 |
+
img_patch = img_patch.transpose([0, 2, 1]).reshape([-1, 3, patch_size, patch_size])
|
143 |
+
result_patch = result_patch.transpose([0, 2, 1]).reshape([-1, 3, patch_size, patch_size])
|
144 |
+
#*----- Stroke Predictor -----*#
|
145 |
+
shape_param, stroke_decision = net_g(img_patch, result_patch)
|
146 |
+
stroke_decision = (stroke_decision > 0).astype('float32')
|
147 |
+
#*----- sampling color -----*#
|
148 |
+
grid = shape_param[:, :, :2].reshape([img_patch.shape[0] * stroke_num, 1, 1, 2])
|
149 |
+
img_temp = img_patch.unsqueeze(1).tile([1, stroke_num, 1, 1, 1]).reshape([
|
150 |
+
img_patch.shape[0] * stroke_num, 3, patch_size, patch_size])
|
151 |
+
color = nn.functional.grid_sample(img_temp, 2 * grid - 1, align_corners=False).reshape([
|
152 |
+
img_patch.shape[0], stroke_num, 3])
|
153 |
+
stroke_param = paddle.concat([shape_param, color], axis=-1)
|
154 |
+
|
155 |
+
param = stroke_param.reshape([-1, 8])
|
156 |
+
decision = stroke_decision.reshape([-1]).astype('bool')
|
157 |
+
param[:, :2] = param[:, :2] / 1.25 + 0.1
|
158 |
+
param[:, 2:4] = param[:, 2:4] / 1.25
|
159 |
+
return param, decision
|
160 |
+
|
161 |
+
|
162 |
+
def sort_strokes(params, decision, scores):
|
163 |
+
"""
|
164 |
+
sort_strokes
|
165 |
+
"""
|
166 |
+
sorted_scores, sorted_index = paddle.sort(scores, axis=1, descending=False)
|
167 |
+
sorted_params = []
|
168 |
+
for idx in range(8):
|
169 |
+
tmp_pick_params = paddle.gather(params[:, :, idx], axis=1, index=sorted_index)
|
170 |
+
sorted_params.append(tmp_pick_params)
|
171 |
+
sorted_params = paddle.stack(sorted_params, axis=2)
|
172 |
+
sorted_decison = paddle.gather(decision.squeeze(2), axis=1, index=sorted_index)
|
173 |
+
return sorted_params, sorted_decison
|
174 |
+
|
175 |
+
|
176 |
+
def render_serial(original_img, net_g, meta_brushes):
|
177 |
+
|
178 |
+
patch_size = 32
|
179 |
+
stroke_num = 8
|
180 |
+
H, W = original_img.shape[-2:]
|
181 |
+
K = max(math.ceil(math.log2(max(H, W) / patch_size)), 0)
|
182 |
+
|
183 |
+
dilation = render_utils.Dilation2d(m=1)
|
184 |
+
erosion = render_utils.Erosion2d(m=1)
|
185 |
+
frames_per_layer = [20, 20, 30, 40, 60]
|
186 |
+
final_frame_list = []
|
187 |
+
|
188 |
+
with paddle.no_grad():
|
189 |
+
#* ----- read in image and init canvas ----- *#
|
190 |
+
final_result = paddle.zeros_like(original_img)
|
191 |
+
|
192 |
+
for layer in range(0, K + 1):
|
193 |
+
t0 = time.time()
|
194 |
+
layer_size = patch_size * (2 ** layer)
|
195 |
+
|
196 |
+
img = nn.functional.interpolate(original_img, (layer_size, layer_size))
|
197 |
+
result = nn.functional.interpolate(final_result, (layer_size, layer_size))
|
198 |
+
img_patch = nn.functional.unfold(img, [patch_size, patch_size],
|
199 |
+
strides=[patch_size, patch_size])
|
200 |
+
result_patch = nn.functional.unfold(result, [patch_size, patch_size],
|
201 |
+
strides=[patch_size, patch_size])
|
202 |
+
h = (img.shape[2] - patch_size) // patch_size + 1
|
203 |
+
w = (img.shape[3] - patch_size) // patch_size + 1
|
204 |
+
render_size_y = int(1.25 * H // h)
|
205 |
+
render_size_x = int(1.25 * W // w)
|
206 |
+
|
207 |
+
#* -------------------------------------------------------------*#
|
208 |
+
#* -------------generate strokes on window type A---------------*#
|
209 |
+
#* -------------------------------------------------------------*#
|
210 |
+
param, decision = stroke_net_predict(img_patch, result_patch, patch_size, net_g, stroke_num)
|
211 |
+
expand_img = original_img
|
212 |
+
wA_xid_list, wA_yid_list, wA_fore_list, wA_alpha_list, wA_error_list, wA_params = \
|
213 |
+
get_single_layer_lists(param, decision, original_img, render_size_x, render_size_y, h, w,
|
214 |
+
meta_brushes, dilation, erosion, stroke_num)
|
215 |
+
|
216 |
+
#* -------------------------------------------------------------*#
|
217 |
+
#* -------------generate strokes on window type B---------------*#
|
218 |
+
#* -------------------------------------------------------------*#
|
219 |
+
#*----- generate input canvas and target patches -----*#
|
220 |
+
wB_error_list = []
|
221 |
+
|
222 |
+
img = nn.functional.pad(img, [patch_size // 2, patch_size // 2,
|
223 |
+
patch_size // 2, patch_size // 2])
|
224 |
+
result = nn.functional.pad(result, [patch_size // 2, patch_size // 2,
|
225 |
+
patch_size // 2, patch_size // 2])
|
226 |
+
img_patch = nn.functional.unfold(img, [patch_size, patch_size],
|
227 |
+
strides=[patch_size, patch_size])
|
228 |
+
result_patch = nn.functional.unfold(result, [patch_size, patch_size],
|
229 |
+
strides=[patch_size, patch_size])
|
230 |
+
h += 1
|
231 |
+
w += 1
|
232 |
+
|
233 |
+
param, decision = stroke_net_predict(img_patch, result_patch, patch_size, net_g, stroke_num)
|
234 |
+
|
235 |
+
patch_y = 4 * render_size_y // 5
|
236 |
+
patch_x = 4 * render_size_x // 5
|
237 |
+
expand_img = nn.functional.pad(original_img, [patch_x // 2, patch_x // 2,
|
238 |
+
patch_y // 2, patch_y // 2])
|
239 |
+
wB_xid_list, wB_yid_list, wB_fore_list, wB_alpha_list, wB_error_list, wB_params = \
|
240 |
+
get_single_layer_lists(param, decision, expand_img, render_size_x, render_size_y, h, w,
|
241 |
+
meta_brushes, dilation, erosion, stroke_num)
|
242 |
+
#* -------------------------------------------------------------*#
|
243 |
+
#* -------------rank strokes and plot stroke one by one---------*#
|
244 |
+
#* -------------------------------------------------------------*#
|
245 |
+
numA = len(wA_error_list)
|
246 |
+
numB = len(wB_error_list)
|
247 |
+
total_error_list = wA_error_list + wB_error_list
|
248 |
+
sort_list = list(np.argsort(total_error_list))
|
249 |
+
|
250 |
+
sample = 0
|
251 |
+
samples = np.linspace(0, len(sort_list) - 2, frames_per_layer[layer]).astype(int)
|
252 |
+
for ii in sort_list:
|
253 |
+
ii = int(ii)
|
254 |
+
if ii < numA:
|
255 |
+
x_id = wA_xid_list[ii]
|
256 |
+
y_id = wA_yid_list[ii]
|
257 |
+
valid_foregrounds = wA_fore_list[ii]
|
258 |
+
valid_alphas = wA_alpha_list[ii]
|
259 |
+
sparam = wA_params[ii]
|
260 |
+
tmp_foreground, tmp_alpha = get_single_stroke_on_full_image_A(x_id, y_id,
|
261 |
+
valid_foregrounds, valid_alphas, sparam, original_img, render_size_x, render_size_y, patch_x, patch_y)
|
262 |
+
else:
|
263 |
+
x_id = wB_xid_list[ii - numA]
|
264 |
+
y_id = wB_yid_list[ii - numA]
|
265 |
+
valid_foregrounds = wB_fore_list[ii - numA]
|
266 |
+
valid_alphas = wB_alpha_list[ii - numA]
|
267 |
+
sparam = wB_params[ii - numA]
|
268 |
+
tmp_foreground, tmp_alpha = get_single_stroke_on_full_image_B(x_id, y_id,
|
269 |
+
valid_foregrounds, valid_alphas, sparam, original_img, render_size_x, render_size_y, patch_x, patch_y)
|
270 |
+
|
271 |
+
final_result = tmp_foreground * tmp_alpha + (1 - tmp_alpha) * final_result
|
272 |
+
if sample in samples:
|
273 |
+
saveframe = (final_result.numpy().squeeze().transpose([1,2,0])[:,:,::-1] * 255).astype(np.uint8)
|
274 |
+
final_frame_list.append(saveframe)
|
275 |
+
#saveframe = cv2.resize(saveframe, (ow, oh))
|
276 |
+
|
277 |
+
sample += 1
|
278 |
+
print("layer %d cost: %.02f" %(layer, time.time() - t0))
|
279 |
+
|
280 |
+
|
281 |
+
saveframe = (final_result.numpy().squeeze().transpose([1,2,0])[:,:,::-1] * 255).astype(np.uint8)
|
282 |
+
final_frame_list.append(saveframe)
|
283 |
+
return final_frame_list
|
render_utils.py
ADDED
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import paddle
|
2 |
+
import paddle.nn as nn
|
3 |
+
import paddle.nn.functional as F
|
4 |
+
import cv2
|
5 |
+
import numpy as np
|
6 |
+
from PIL import Image
|
7 |
+
import math
|
8 |
+
|
9 |
+
class Erosion2d(nn.Layer):
|
10 |
+
"""
|
11 |
+
Erosion2d
|
12 |
+
"""
|
13 |
+
def __init__(self, m=1):
|
14 |
+
super(Erosion2d, self).__init__()
|
15 |
+
self.m = m
|
16 |
+
self.pad = [m, m, m, m]
|
17 |
+
|
18 |
+
def forward(self, x):
|
19 |
+
batch_size, c, h, w = x.shape
|
20 |
+
x_pad = F.pad(x, pad=self.pad, mode='constant', value=1e9)
|
21 |
+
channel = nn.functional.unfold(x_pad, 2 * self.m + 1, strides=1, paddings=0).reshape([batch_size, c, -1, h, w])
|
22 |
+
result = paddle.min(channel, axis=2)
|
23 |
+
return result
|
24 |
+
|
25 |
+
class Dilation2d(nn.Layer):
|
26 |
+
"""
|
27 |
+
Dilation2d
|
28 |
+
"""
|
29 |
+
def __init__(self, m=1):
|
30 |
+
super(Dilation2d, self).__init__()
|
31 |
+
self.m = m
|
32 |
+
self.pad = [m, m, m, m]
|
33 |
+
|
34 |
+
def forward(self, x):
|
35 |
+
batch_size, c, h, w = x.shape
|
36 |
+
x_pad = F.pad(x, pad=self.pad, mode='constant', value=-1e9)
|
37 |
+
channel = nn.functional.unfold(x_pad, 2 * self.m + 1, strides=1, paddings=0).reshape([batch_size, c, -1, h, w])
|
38 |
+
result = paddle.max(channel, axis=2)
|
39 |
+
return result
|
40 |
+
|
41 |
+
def param2stroke(param, H, W, meta_brushes):
|
42 |
+
"""
|
43 |
+
param2stroke
|
44 |
+
"""
|
45 |
+
b = param.shape[0]
|
46 |
+
param_list = paddle.split(param, 8, axis=1)
|
47 |
+
x0, y0, w, h, theta = [item.squeeze(-1) for item in param_list[:5]]
|
48 |
+
sin_theta = paddle.sin(math.pi * theta)
|
49 |
+
cos_theta = paddle.cos(math.pi * theta)
|
50 |
+
index = paddle.full((b,), -1, dtype='int64').numpy()
|
51 |
+
|
52 |
+
index[(h > w).numpy()] = 0
|
53 |
+
index[(h <= w).numpy()] = 1
|
54 |
+
meta_brushes_resize = F.interpolate(meta_brushes, (H, W)).numpy()
|
55 |
+
brush = paddle.to_tensor(meta_brushes_resize[index])
|
56 |
+
|
57 |
+
warp_00 = cos_theta / w
|
58 |
+
warp_01 = sin_theta * H / (W * w)
|
59 |
+
warp_02 = (1 - 2 * x0) * cos_theta / w + (1 - 2 * y0) * sin_theta * H / (W * w)
|
60 |
+
warp_10 = -sin_theta * W / (H * h)
|
61 |
+
warp_11 = cos_theta / h
|
62 |
+
warp_12 = (1 - 2 * y0) * cos_theta / h - (1 - 2 * x0) * sin_theta * W / (H * h)
|
63 |
+
warp_0 = paddle.stack([warp_00, warp_01, warp_02], axis=1)
|
64 |
+
warp_1 = paddle.stack([warp_10, warp_11, warp_12], axis=1)
|
65 |
+
warp = paddle.stack([warp_0, warp_1], axis=1)
|
66 |
+
grid = nn.functional.affine_grid(warp, [b, 3, H, W]) # paddle和torch默认值是反过来的
|
67 |
+
brush = nn.functional.grid_sample(brush, grid)
|
68 |
+
return brush
|
69 |
+
|
70 |
+
|
71 |
+
def read_img(img_path, img_type='RGB', h=None, w=None):
|
72 |
+
"""
|
73 |
+
read img
|
74 |
+
"""
|
75 |
+
img = Image.open(img_path).convert(img_type)
|
76 |
+
if h is not None and w is not None:
|
77 |
+
img = img.resize((w, h), resample=Image.NEAREST)
|
78 |
+
img = np.array(img)
|
79 |
+
if img.ndim == 2:
|
80 |
+
img = np.expand_dims(img, axis=-1)
|
81 |
+
img = img.transpose((2, 0, 1))
|
82 |
+
img = paddle.to_tensor(img).unsqueeze(0).astype('float32') / 255.
|
83 |
+
return img
|
84 |
+
|
85 |
+
def preprocess(img, w=512, h=512):
|
86 |
+
image = cv2.resize(img, (w, h), cv2.INTER_NEAREST)
|
87 |
+
image = image.transpose((2, 0, 1))
|
88 |
+
image = paddle.to_tensor(image).unsqueeze(0).astype('float32') / 255.
|
89 |
+
return image
|
90 |
+
|
91 |
+
def pad(img, H, W):
|
92 |
+
b, c, h, w = img.shape
|
93 |
+
pad_h = (H - h) // 2
|
94 |
+
pad_w = (W - w) // 2
|
95 |
+
remainder_h = (H - h) % 2
|
96 |
+
remainder_w = (W - w) % 2
|
97 |
+
expand_img = nn.functional.pad(img, [pad_w, pad_w + remainder_w,
|
98 |
+
pad_h, pad_h + remainder_h])
|
99 |
+
return expand_img
|
100 |
+
|
101 |
+
|
102 |
+
|
requirements.txt
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
gdown
|
2 |
+
numpy
|
3 |
+
loguru
|
4 |
+
torch
|
5 |
+
gradio
|
6 |
+
Pillow
|
7 |
+
opencv-python
|
8 |
+
paddlepaddle==2.2.1
|