jaekookang commited on
Commit
9ff1108
1 Parent(s): a871eb8

first upload

Browse files
.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: app.py
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