Spaces:
Runtime error
Runtime error
Xu Ma
commited on
Commit
•
c89c010
1
Parent(s):
1b90f20
update
Browse files- app.py +207 -138
- config/base.yaml +19 -3
- main.py +339 -0
app.py
CHANGED
@@ -9,124 +9,193 @@ import torch
|
|
9 |
import yaml
|
10 |
from PIL import Image
|
11 |
from subprocess import call
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
)
|
28 |
-
parser.add_argument(
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
)
|
35 |
-
parser.add_argument(
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
)
|
42 |
-
parser.add_argument(
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
)
|
49 |
-
parser.
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
130 |
|
131 |
|
132 |
def run_cmd(command):
|
@@ -150,25 +219,25 @@ run_cmd("python main.py --config config/base.yaml --experiment experiment_5x1 --
|
|
150 |
|
151 |
|
152 |
|
153 |
-
# yaml文件解析
|
154 |
-
def yaml_parse(file_path):
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
# yaml csv 文件解析
|
159 |
-
def yaml_csv(file_path, file_tag):
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
|
173 |
|
174 |
def main(args):
|
@@ -223,7 +292,7 @@ def main(args):
|
|
223 |
|
224 |
# Interface
|
225 |
gr.Interface(
|
226 |
-
fn=
|
227 |
inputs=inputs,
|
228 |
outputs=[outputs, outputs02],
|
229 |
title=title,
|
|
|
9 |
import yaml
|
10 |
from PIL import Image
|
11 |
from subprocess import call
|
12 |
+
import pydiffvg
|
13 |
+
import torch
|
14 |
+
import cv2
|
15 |
+
import matplotlib.pyplot as plt
|
16 |
+
import random
|
17 |
+
import argparse
|
18 |
+
import math
|
19 |
+
import errno
|
20 |
+
from tqdm import tqdm
|
21 |
+
import yaml
|
22 |
+
from easydict import EasyDict as edict
|
23 |
+
from main import main_func
|
24 |
+
|
25 |
+
def parse_args():
|
26 |
+
parser = argparse.ArgumentParser()
|
27 |
+
parser.add_argument('--debug', action='store_true', default=False)
|
28 |
+
parser.add_argument("--config", default="config/base.yaml", type=str)
|
29 |
+
parser.add_argument("--experiment", type=str)
|
30 |
+
parser.add_argument("--seed", type=int)
|
31 |
+
parser.add_argument("--target", type=str, help="target image path")
|
32 |
+
parser.add_argument('--log_dir', metavar='DIR', default="log/debug")
|
33 |
+
parser.add_argument('--initial', type=str, default="random", choices=['random', 'circle'])
|
34 |
+
parser.add_argument('--signature', nargs='+', type=str)
|
35 |
+
parser.add_argument('--seginit', nargs='+', type=str)
|
36 |
+
parser.add_argument("--num_segments", type=int, default=4)
|
37 |
+
# parser.add_argument("--num_paths", type=str, default="1,1,1")
|
38 |
+
# parser.add_argument("--num_iter", type=int, default=500)
|
39 |
+
# parser.add_argument('--free', action='store_true')
|
40 |
+
# Please ensure that image resolution is divisible by pool_size; otherwise the performance would drop a lot.
|
41 |
+
# parser.add_argument('--pool_size', type=int, default=40, help="the pooled image size for next path initialization")
|
42 |
+
# parser.add_argument('--save_loss', action='store_true')
|
43 |
+
# parser.add_argument('--save_init', action='store_true')
|
44 |
+
# parser.add_argument('--save_image', action='store_true')
|
45 |
+
# parser.add_argument('--save_video', action='store_true')
|
46 |
+
# parser.add_argument('--print_weight', action='store_true')
|
47 |
+
# parser.add_argument('--circle_init_radius', type=float)
|
48 |
+
cfg = edict()
|
49 |
+
args = parser.parse_args()
|
50 |
+
cfg.debug = args.debug
|
51 |
+
cfg.config = args.config
|
52 |
+
cfg.experiment = args.experiment
|
53 |
+
cfg.seed = args.seed
|
54 |
+
cfg.target = args.target
|
55 |
+
cfg.log_dir = args.log_dir
|
56 |
+
cfg.initial = args.initial
|
57 |
+
cfg.signature = args.signature
|
58 |
+
# set cfg num_segments in command
|
59 |
+
cfg.num_segments = args.num_segments
|
60 |
+
if args.seginit is not None:
|
61 |
+
cfg.seginit = edict()
|
62 |
+
cfg.seginit.type = args.seginit[0]
|
63 |
+
if cfg.seginit.type == 'circle':
|
64 |
+
cfg.seginit.radius = float(args.seginit[1])
|
65 |
+
return cfg
|
66 |
+
|
67 |
+
|
68 |
+
|
69 |
+
|
70 |
+
|
71 |
+
def run_live(img, experiment_id):
|
72 |
+
main_func(img, experiment_id)
|
73 |
+
return 0, 1
|
74 |
+
|
75 |
+
|
76 |
+
|
77 |
+
|
78 |
+
|
79 |
+
|
80 |
+
|
81 |
+
|
82 |
+
|
83 |
+
# ROOT_PATH = sys.path[0] # 根目录
|
84 |
+
# # 模型路径
|
85 |
+
# model_path = "ultralytics/yolov5"
|
86 |
+
# # 模型名称临时变量
|
87 |
+
# model_name_tmp = ""
|
88 |
+
# # 设备临时变量
|
89 |
+
# device_tmp = ""
|
90 |
+
# # 文件后缀
|
91 |
+
# suffix_list = [".csv", ".yaml"]
|
92 |
+
# def parse_args(known=False):
|
93 |
+
# parser = argparse.ArgumentParser(description="Gradio LIVE")
|
94 |
+
# parser.add_argument(
|
95 |
+
# "--model_name", "-mn", default="yolov5s", type=str, help="model name"
|
96 |
+
# )
|
97 |
+
# parser.add_argument(
|
98 |
+
# "--model_cfg",
|
99 |
+
# "-mc",
|
100 |
+
# default="./model_config/model_name_p5_all.yaml",
|
101 |
+
# type=str,
|
102 |
+
# help="model config",
|
103 |
+
# )
|
104 |
+
# parser.add_argument(
|
105 |
+
# "--cls_name",
|
106 |
+
# "-cls",
|
107 |
+
# default="./cls_name/cls_name.yaml",
|
108 |
+
# type=str,
|
109 |
+
# help="cls name",
|
110 |
+
# )
|
111 |
+
# parser.add_argument(
|
112 |
+
# "--nms_conf",
|
113 |
+
# "-conf",
|
114 |
+
# default=0.5,
|
115 |
+
# type=float,
|
116 |
+
# help="model NMS confidence threshold",
|
117 |
+
# )
|
118 |
+
# parser.add_argument(
|
119 |
+
# "--nms_iou", "-iou", default=0.45, type=float, help="model NMS IoU threshold"
|
120 |
+
# )
|
121 |
+
#
|
122 |
+
# parser.add_argument(
|
123 |
+
# "--label_dnt_show",
|
124 |
+
# "-lds",
|
125 |
+
# action="store_false",
|
126 |
+
# default=True,
|
127 |
+
# help="label show",
|
128 |
+
# )
|
129 |
+
# parser.add_argument(
|
130 |
+
# "--device",
|
131 |
+
# "-dev",
|
132 |
+
# default="cpu",
|
133 |
+
# type=str,
|
134 |
+
# help="cuda or cpu, hugging face only cpu",
|
135 |
+
# )
|
136 |
+
# parser.add_argument(
|
137 |
+
# "--inference_size", "-isz", default=640, type=int, help="model inference size"
|
138 |
+
# )
|
139 |
+
#
|
140 |
+
# args = parser.parse_known_args()[0] if known else parser.parse_args()
|
141 |
+
# return args
|
142 |
+
# # 模型加载
|
143 |
+
# def model_loading(model_name, device):
|
144 |
+
#
|
145 |
+
# # 加载本地模型
|
146 |
+
# model = torch.hub.load(model_path, model_name, force_reload=True, device=device)
|
147 |
+
#
|
148 |
+
# return model
|
149 |
+
# # 检测信息
|
150 |
+
# def export_json(results, model, img_size):
|
151 |
+
#
|
152 |
+
# return [
|
153 |
+
# [
|
154 |
+
# {
|
155 |
+
# "id": int(i),
|
156 |
+
# "class": int(result[i][5]),
|
157 |
+
# "class_name": model.model.names[int(result[i][5])],
|
158 |
+
# "normalized_box": {
|
159 |
+
# "x0": round(result[i][:4].tolist()[0], 6),
|
160 |
+
# "y0": round(result[i][:4].tolist()[1], 6),
|
161 |
+
# "x1": round(result[i][:4].tolist()[2], 6),
|
162 |
+
# "y1": round(result[i][:4].tolist()[3], 6),
|
163 |
+
# },
|
164 |
+
# "confidence": round(float(result[i][4]), 2),
|
165 |
+
# "fps": round(1000 / float(results.t[1]), 2),
|
166 |
+
# "width": img_size[0],
|
167 |
+
# "height": img_size[1],
|
168 |
+
# }
|
169 |
+
# for i in range(len(result))
|
170 |
+
# ]
|
171 |
+
# for result in results.xyxyn
|
172 |
+
# ]
|
173 |
+
# def yolo_det(img, experiment_id, device=None, model_name=None, inference_size=None, conf=None, iou=None, label_opt=None, model_cls=None):
|
174 |
+
#
|
175 |
+
# global model, model_name_tmp, device_tmp
|
176 |
+
#
|
177 |
+
# if model_name_tmp != model_name:
|
178 |
+
# # 模型判断,避免反复加载
|
179 |
+
# model_name_tmp = model_name
|
180 |
+
# model = model_loading(model_name_tmp, device)
|
181 |
+
# elif device_tmp != device:
|
182 |
+
# device_tmp = device
|
183 |
+
# model = model_loading(model_name_tmp, device)
|
184 |
+
#
|
185 |
+
# # -----------模型调参-----------
|
186 |
+
# model.conf = conf # NMS 置信度阈值
|
187 |
+
# model.iou = iou # NMS IOU阈值
|
188 |
+
# model.max_det = 1000 # 最大检测框数
|
189 |
+
# model.classes = model_cls # 模型类别
|
190 |
+
#
|
191 |
+
# results = model(img, size=inference_size) # 检测
|
192 |
+
# results.render(labels=label_opt) # 渲染
|
193 |
+
#
|
194 |
+
# det_img = Image.fromarray(results.imgs[0]) # 检测图片
|
195 |
+
#
|
196 |
+
# det_json = export_json(results, model, img.size)[0] # 检测信息
|
197 |
+
#
|
198 |
+
# return det_img, det_json
|
199 |
|
200 |
|
201 |
def run_cmd(command):
|
|
|
219 |
|
220 |
|
221 |
|
222 |
+
# # yaml文件解析
|
223 |
+
# def yaml_parse(file_path):
|
224 |
+
# return yaml.safe_load(open(file_path, "r", encoding="utf-8").read())
|
225 |
+
#
|
226 |
+
#
|
227 |
+
# # yaml csv 文件解析
|
228 |
+
# def yaml_csv(file_path, file_tag):
|
229 |
+
# file_suffix = Path(file_path).suffix
|
230 |
+
# if file_suffix == suffix_list[0]:
|
231 |
+
# # 模型名称
|
232 |
+
# file_names = [i[0] for i in list(csv.reader(open(file_path)))] # csv版
|
233 |
+
# elif file_suffix == suffix_list[1]:
|
234 |
+
# # 模型名称
|
235 |
+
# file_names = yaml_parse(file_path).get(file_tag) # yaml版
|
236 |
+
# else:
|
237 |
+
# print(f"{file_path}格式不正确!程序退出!")
|
238 |
+
# sys.exit()
|
239 |
+
#
|
240 |
+
# return file_names
|
241 |
|
242 |
|
243 |
def main(args):
|
|
|
292 |
|
293 |
# Interface
|
294 |
gr.Interface(
|
295 |
+
fn=run_live,
|
296 |
inputs=inputs,
|
297 |
outputs=[outputs, outputs02],
|
298 |
title=title,
|
config/base.yaml
CHANGED
@@ -5,10 +5,10 @@ default:
|
|
5 |
type: circle
|
6 |
radius: 5
|
7 |
save:
|
8 |
-
init:
|
9 |
-
image:
|
10 |
output: true
|
11 |
-
video:
|
12 |
loss: false
|
13 |
trainable:
|
14 |
bg: False
|
@@ -66,3 +66,19 @@ experiment_1357:
|
|
66 |
type: list
|
67 |
schedule: [1, 3, 5, 7]
|
68 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
5 |
type: circle
|
6 |
radius: 5
|
7 |
save:
|
8 |
+
init: false
|
9 |
+
image: false
|
10 |
output: true
|
11 |
+
video: false
|
12 |
loss: false
|
13 |
trainable:
|
14 |
bg: False
|
|
|
66 |
type: list
|
67 |
schedule: [1, 3, 5, 7]
|
68 |
|
69 |
+
|
70 |
+
experiment_exp2_256:
|
71 |
+
path_schedule:
|
72 |
+
type: exp
|
73 |
+
base: 2
|
74 |
+
max_path: 256
|
75 |
+
max_path_per_iter: 32
|
76 |
+
|
77 |
+
|
78 |
+
experiment_exp2_128:
|
79 |
+
path_schedule:
|
80 |
+
type: exp
|
81 |
+
base: 2
|
82 |
+
max_path: 128
|
83 |
+
max_path_per_iter: 32
|
84 |
+
|
main.py
CHANGED
@@ -344,6 +344,345 @@ class linear_decay_lrlambda_f(object):
|
|
344 |
lr = lr_s * (1-r) + lr_e * r
|
345 |
return lr
|
346 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
347 |
|
348 |
if __name__ == "__main__":
|
349 |
|
|
|
344 |
lr = lr_s * (1-r) + lr_e * r
|
345 |
return lr
|
346 |
|
347 |
+
def main_func(target, experiment):
|
348 |
+
cfg_arg = parse_args()
|
349 |
+
with open(cfg_arg.config, 'r') as f:
|
350 |
+
cfg = yaml.load(f, Loader=yaml.FullLoader)
|
351 |
+
cfg_default = edict(cfg['default'])
|
352 |
+
cfg = edict(cfg[cfg_arg.experiment])
|
353 |
+
cfg.update(cfg_default)
|
354 |
+
cfg.update(cfg_arg)
|
355 |
+
cfg.exid = get_experiment_id(cfg.debug)
|
356 |
+
|
357 |
+
cfg.experiment_dir = \
|
358 |
+
osp.join(cfg.log_dir, '{}_{}'.format(cfg.exid, '_'.join(cfg.signature)))
|
359 |
+
cfg.target = target
|
360 |
+
cfg.experiment = experiment
|
361 |
+
|
362 |
+
configfile = osp.join(cfg.experiment_dir, 'config.yaml')
|
363 |
+
check_and_create_dir(configfile)
|
364 |
+
with open(osp.join(configfile), 'w') as f:
|
365 |
+
yaml.dump(edict_2_dict(cfg), f)
|
366 |
+
|
367 |
+
# Use GPU if available
|
368 |
+
pydiffvg.set_use_gpu(torch.cuda.is_available())
|
369 |
+
device = pydiffvg.get_device()
|
370 |
+
|
371 |
+
gt = np.array(PIL.Image.open(cfg.target))
|
372 |
+
print(f"Input image shape is: {gt.shape}")
|
373 |
+
if len(gt.shape) == 2:
|
374 |
+
print("Converting the gray-scale image to RGB.")
|
375 |
+
gt = gt.unsqueeze(dim=-1).repeat(1,1,3)
|
376 |
+
if gt.shape[2] == 4:
|
377 |
+
print("Input image includes alpha channel, simply dropout alpha channel.")
|
378 |
+
gt = gt[:, :, :3]
|
379 |
+
gt = (gt/255).astype(np.float32)
|
380 |
+
gt = torch.FloatTensor(gt).permute(2, 0, 1)[None].to(device)
|
381 |
+
if cfg.use_ycrcb:
|
382 |
+
gt = ycrcb_conversion(gt)
|
383 |
+
h, w = gt.shape[2:]
|
384 |
+
|
385 |
+
path_schedule = get_path_schedule(**cfg.path_schedule)
|
386 |
+
|
387 |
+
if cfg.seed is not None:
|
388 |
+
random.seed(cfg.seed)
|
389 |
+
npr.seed(cfg.seed)
|
390 |
+
torch.manual_seed(cfg.seed)
|
391 |
+
render = pydiffvg.RenderFunction.apply
|
392 |
+
|
393 |
+
shapes_record, shape_groups_record = [], []
|
394 |
+
|
395 |
+
region_loss = None
|
396 |
+
loss_matrix = []
|
397 |
+
|
398 |
+
para_point, para_color = {}, {}
|
399 |
+
if cfg.trainable.stroke:
|
400 |
+
para_stroke_width, para_stroke_color = {}, {}
|
401 |
+
|
402 |
+
pathn_record = []
|
403 |
+
# Background
|
404 |
+
if cfg.trainable.bg:
|
405 |
+
# meancolor = gt.mean([2, 3])[0]
|
406 |
+
para_bg = torch.tensor([1., 1., 1.], requires_grad=True, device=device)
|
407 |
+
else:
|
408 |
+
if cfg.use_ycrcb:
|
409 |
+
para_bg = torch.tensor([219/255, 0, 0], requires_grad=False, device=device)
|
410 |
+
else:
|
411 |
+
para_bg = torch.tensor([1., 1., 1.], requires_grad=False, device=device)
|
412 |
+
|
413 |
+
##################
|
414 |
+
# start_training #
|
415 |
+
##################
|
416 |
+
|
417 |
+
loss_weight = None
|
418 |
+
loss_weight_keep = 0
|
419 |
+
if cfg.coord_init.type == 'naive':
|
420 |
+
pos_init_method = naive_coord_init(
|
421 |
+
para_bg.view(1, -1, 1, 1).repeat(1, 1, h, w), gt)
|
422 |
+
elif cfg.coord_init.type == 'sparse':
|
423 |
+
pos_init_method = sparse_coord_init(
|
424 |
+
para_bg.view(1, -1, 1, 1).repeat(1, 1, h, w), gt)
|
425 |
+
elif cfg.coord_init.type == 'random':
|
426 |
+
pos_init_method = random_coord_init([h, w])
|
427 |
+
else:
|
428 |
+
raise ValueError
|
429 |
+
|
430 |
+
lrlambda_f = linear_decay_lrlambda_f(cfg.num_iter, 0.4)
|
431 |
+
optim_schedular_dict = {}
|
432 |
+
|
433 |
+
for path_idx, pathn in enumerate(path_schedule):
|
434 |
+
loss_list = []
|
435 |
+
print("=> Adding [{}] paths, [{}] ...".format(pathn, cfg.seginit.type))
|
436 |
+
pathn_record.append(pathn)
|
437 |
+
pathn_record_str = '-'.join([str(i) for i in pathn_record])
|
438 |
+
|
439 |
+
# initialize new shapes related stuffs.
|
440 |
+
if cfg.trainable.stroke:
|
441 |
+
shapes, shape_groups, point_var, color_var, stroke_width_var, stroke_color_var = init_shapes(
|
442 |
+
pathn, cfg.num_segments, (h, w),
|
443 |
+
cfg.seginit, len(shapes_record),
|
444 |
+
pos_init_method,
|
445 |
+
trainable_stroke=True,
|
446 |
+
gt=gt, )
|
447 |
+
para_stroke_width[path_idx] = stroke_width_var
|
448 |
+
para_stroke_color[path_idx] = stroke_color_var
|
449 |
+
else:
|
450 |
+
shapes, shape_groups, point_var, color_var = init_shapes(
|
451 |
+
pathn, cfg.num_segments, (h, w),
|
452 |
+
cfg.seginit, len(shapes_record),
|
453 |
+
pos_init_method,
|
454 |
+
trainable_stroke=False,
|
455 |
+
gt=gt, )
|
456 |
+
|
457 |
+
shapes_record += shapes
|
458 |
+
shape_groups_record += shape_groups
|
459 |
+
|
460 |
+
if cfg.save.init:
|
461 |
+
filename = os.path.join(
|
462 |
+
cfg.experiment_dir, "svg-init",
|
463 |
+
"{}-init.svg".format(pathn_record_str))
|
464 |
+
check_and_create_dir(filename)
|
465 |
+
pydiffvg.save_svg(
|
466 |
+
filename, w, h,
|
467 |
+
shapes_record, shape_groups_record)
|
468 |
+
|
469 |
+
para = {}
|
470 |
+
if (cfg.trainable.bg) and (path_idx == 0):
|
471 |
+
para['bg'] = [para_bg]
|
472 |
+
para['point'] = point_var
|
473 |
+
para['color'] = color_var
|
474 |
+
if cfg.trainable.stroke:
|
475 |
+
para['stroke_width'] = stroke_width_var
|
476 |
+
para['stroke_color'] = stroke_color_var
|
477 |
+
|
478 |
+
pg = [{'params' : para[ki], 'lr' : cfg.lr_base[ki]} for ki in sorted(para.keys())]
|
479 |
+
optim = torch.optim.Adam(pg)
|
480 |
+
|
481 |
+
if cfg.trainable.record:
|
482 |
+
scheduler = LambdaLR(
|
483 |
+
optim, lr_lambda=lrlambda_f, last_epoch=-1)
|
484 |
+
else:
|
485 |
+
scheduler = LambdaLR(
|
486 |
+
optim, lr_lambda=lrlambda_f, last_epoch=cfg.num_iter)
|
487 |
+
optim_schedular_dict[path_idx] = (optim, scheduler)
|
488 |
+
|
489 |
+
# Inner loop training
|
490 |
+
t_range = tqdm(range(cfg.num_iter))
|
491 |
+
for t in t_range:
|
492 |
+
|
493 |
+
for _, (optim, _) in optim_schedular_dict.items():
|
494 |
+
optim.zero_grad()
|
495 |
+
|
496 |
+
# Forward pass: render the image.
|
497 |
+
scene_args = pydiffvg.RenderFunction.serialize_scene(
|
498 |
+
w, h, shapes_record, shape_groups_record)
|
499 |
+
img = render(w, h, 2, 2, t, None, *scene_args)
|
500 |
+
|
501 |
+
# Compose img with white background
|
502 |
+
img = img[:, :, 3:4] * img[:, :, :3] + \
|
503 |
+
para_bg * (1 - img[:, :, 3:4])
|
504 |
+
|
505 |
+
if cfg.save.video:
|
506 |
+
filename = os.path.join(
|
507 |
+
cfg.experiment_dir, "video-png",
|
508 |
+
"{}-iter{}.png".format(pathn_record_str, t))
|
509 |
+
check_and_create_dir(filename)
|
510 |
+
if cfg.use_ycrcb:
|
511 |
+
imshow = ycrcb_conversion(
|
512 |
+
img, format='[2D x 3]', reverse=True).detach().cpu()
|
513 |
+
else:
|
514 |
+
imshow = img.detach().cpu()
|
515 |
+
pydiffvg.imwrite(imshow, filename, gamma=gamma)
|
516 |
+
|
517 |
+
x = img.unsqueeze(0).permute(0, 3, 1, 2) # HWC -> NCHW
|
518 |
+
|
519 |
+
if cfg.use_ycrcb:
|
520 |
+
color_reweight = torch.FloatTensor([255/219, 255/224, 255/255]).to(device)
|
521 |
+
loss = ((x-gt)*(color_reweight.view(1, -1, 1, 1)))**2
|
522 |
+
else:
|
523 |
+
loss = ((x-gt)**2)
|
524 |
+
|
525 |
+
if cfg.loss.use_l1_loss:
|
526 |
+
loss = abs(x-gt)
|
527 |
+
|
528 |
+
if cfg.loss.use_distance_weighted_loss:
|
529 |
+
if cfg.use_ycrcb:
|
530 |
+
raise ValueError
|
531 |
+
shapes_forsdf = copy.deepcopy(shapes)
|
532 |
+
shape_groups_forsdf = copy.deepcopy(shape_groups)
|
533 |
+
for si in shapes_forsdf:
|
534 |
+
si.stroke_width = torch.FloatTensor([0]).to(device)
|
535 |
+
for sg_idx, sgi in enumerate(shape_groups_forsdf):
|
536 |
+
sgi.fill_color = torch.FloatTensor([1, 1, 1, 1]).to(device)
|
537 |
+
sgi.shape_ids = torch.LongTensor([sg_idx]).to(device)
|
538 |
+
|
539 |
+
sargs_forsdf = pydiffvg.RenderFunction.serialize_scene(
|
540 |
+
w, h, shapes_forsdf, shape_groups_forsdf)
|
541 |
+
with torch.no_grad():
|
542 |
+
im_forsdf = render(w, h, 2, 2, 0, None, *sargs_forsdf)
|
543 |
+
# use alpha channel is a trick to get 0-1 image
|
544 |
+
im_forsdf = (im_forsdf[:, :, 3]).detach().cpu().numpy()
|
545 |
+
loss_weight = get_sdf(im_forsdf, normalize='to1')
|
546 |
+
loss_weight += loss_weight_keep
|
547 |
+
loss_weight = np.clip(loss_weight, 0, 1)
|
548 |
+
loss_weight = torch.FloatTensor(loss_weight).to(device)
|
549 |
+
|
550 |
+
if cfg.save.loss:
|
551 |
+
save_loss = loss.squeeze(dim=0).mean(dim=0,keepdim=False).cpu().detach().numpy()
|
552 |
+
save_weight = loss_weight.cpu().detach().numpy()
|
553 |
+
save_weighted_loss = save_loss*save_weight
|
554 |
+
# normalize to [0,1]
|
555 |
+
save_loss = (save_loss - np.min(save_loss))/np.ptp(save_loss)
|
556 |
+
save_weight = (save_weight - np.min(save_weight))/np.ptp(save_weight)
|
557 |
+
save_weighted_loss = (save_weighted_loss - np.min(save_weighted_loss))/np.ptp(save_weighted_loss)
|
558 |
+
|
559 |
+
# save
|
560 |
+
plt.imshow(save_loss, cmap='Reds')
|
561 |
+
plt.axis('off')
|
562 |
+
# plt.colorbar()
|
563 |
+
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-mseloss.png".format(pathn_record_str, t))
|
564 |
+
check_and_create_dir(filename)
|
565 |
+
plt.savefig(filename, dpi=800)
|
566 |
+
plt.close()
|
567 |
+
|
568 |
+
plt.imshow(save_weight, cmap='Greys')
|
569 |
+
plt.axis('off')
|
570 |
+
# plt.colorbar()
|
571 |
+
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-sdfweight.png".format(pathn_record_str, t))
|
572 |
+
plt.savefig(filename, dpi=800)
|
573 |
+
plt.close()
|
574 |
+
|
575 |
+
plt.imshow(save_weighted_loss, cmap='Reds')
|
576 |
+
plt.axis('off')
|
577 |
+
# plt.colorbar()
|
578 |
+
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-weightedloss.png".format(pathn_record_str, t))
|
579 |
+
plt.savefig(filename, dpi=800)
|
580 |
+
plt.close()
|
581 |
+
|
582 |
+
|
583 |
+
|
584 |
+
|
585 |
+
|
586 |
+
if loss_weight is None:
|
587 |
+
loss = loss.sum(1).mean()
|
588 |
+
else:
|
589 |
+
loss = (loss.sum(1)*loss_weight).mean()
|
590 |
+
|
591 |
+
# if (cfg.loss.bis_loss_weight is not None) and (cfg.loss.bis_loss_weight > 0):
|
592 |
+
# loss_bis = bezier_intersection_loss(point_var[0]) * cfg.loss.bis_loss_weight
|
593 |
+
# loss = loss + loss_bis
|
594 |
+
if (cfg.loss.xing_loss_weight is not None) \
|
595 |
+
and (cfg.loss.xing_loss_weight > 0):
|
596 |
+
loss_xing = xing_loss(point_var) * cfg.loss.xing_loss_weight
|
597 |
+
loss = loss + loss_xing
|
598 |
+
|
599 |
+
|
600 |
+
loss_list.append(loss.item())
|
601 |
+
t_range.set_postfix({'loss': loss.item()})
|
602 |
+
loss.backward()
|
603 |
+
|
604 |
+
# step
|
605 |
+
for _, (optim, scheduler) in optim_schedular_dict.items():
|
606 |
+
optim.step()
|
607 |
+
scheduler.step()
|
608 |
+
|
609 |
+
for group in shape_groups_record:
|
610 |
+
group.fill_color.data.clamp_(0.0, 1.0)
|
611 |
+
|
612 |
+
if cfg.loss.use_distance_weighted_loss:
|
613 |
+
loss_weight_keep = loss_weight.detach().cpu().numpy() * 1
|
614 |
+
|
615 |
+
if not cfg.trainable.record:
|
616 |
+
for _, pi in pg.items():
|
617 |
+
for ppi in pi:
|
618 |
+
pi.require_grad = False
|
619 |
+
optim_schedular_dict = {}
|
620 |
+
|
621 |
+
if cfg.save.image:
|
622 |
+
filename = os.path.join(
|
623 |
+
cfg.experiment_dir, "demo-png", "{}.png".format(pathn_record_str))
|
624 |
+
check_and_create_dir(filename)
|
625 |
+
if cfg.use_ycrcb:
|
626 |
+
imshow = ycrcb_conversion(
|
627 |
+
img, format='[2D x 3]', reverse=True).detach().cpu()
|
628 |
+
else:
|
629 |
+
imshow = img.detach().cpu()
|
630 |
+
pydiffvg.imwrite(imshow, filename, gamma=gamma)
|
631 |
+
|
632 |
+
if cfg.save.output:
|
633 |
+
filename = os.path.join(
|
634 |
+
cfg.experiment_dir, "output-svg", "{}.svg".format(pathn_record_str))
|
635 |
+
check_and_create_dir(filename)
|
636 |
+
pydiffvg.save_svg(filename, w, h, shapes_record, shape_groups_record)
|
637 |
+
|
638 |
+
loss_matrix.append(loss_list)
|
639 |
+
|
640 |
+
# calculate the pixel loss
|
641 |
+
# pixel_loss = ((x-gt)**2).sum(dim=1, keepdim=True).sqrt_() # [N,1,H, W]
|
642 |
+
# region_loss = adaptive_avg_pool2d(pixel_loss, cfg.region_loss_pool_size)
|
643 |
+
# loss_weight = torch.softmax(region_loss.reshape(1, 1, -1), dim=-1)\
|
644 |
+
# .reshape_as(region_loss)
|
645 |
+
|
646 |
+
pos_init_method = naive_coord_init(x, gt)
|
647 |
+
|
648 |
+
if cfg.coord_init.type == 'naive':
|
649 |
+
pos_init_method = naive_coord_init(x, gt)
|
650 |
+
elif cfg.coord_init.type == 'sparse':
|
651 |
+
pos_init_method = sparse_coord_init(x, gt)
|
652 |
+
elif cfg.coord_init.type == 'random':
|
653 |
+
pos_init_method = random_coord_init([h, w])
|
654 |
+
else:
|
655 |
+
raise ValueError
|
656 |
+
|
657 |
+
if cfg.save.video:
|
658 |
+
print("saving iteration video...")
|
659 |
+
img_array = []
|
660 |
+
for ii in range(0, cfg.num_iter):
|
661 |
+
filename = os.path.join(
|
662 |
+
cfg.experiment_dir, "video-png",
|
663 |
+
"{}-iter{}.png".format(pathn_record_str, ii))
|
664 |
+
img = cv2.imread(filename)
|
665 |
+
# cv2.putText(
|
666 |
+
# img, "Path:{} \nIteration:{}".format(pathn_record_str, ii),
|
667 |
+
# (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
|
668 |
+
img_array.append(img)
|
669 |
+
|
670 |
+
videoname = os.path.join(
|
671 |
+
cfg.experiment_dir, "video-avi",
|
672 |
+
"{}.avi".format(pathn_record_str))
|
673 |
+
check_and_create_dir(videoname)
|
674 |
+
out = cv2.VideoWriter(
|
675 |
+
videoname,
|
676 |
+
# cv2.VideoWriter_fourcc(*'mp4v'),
|
677 |
+
cv2.VideoWriter_fourcc(*'FFV1'),
|
678 |
+
20.0, (w, h))
|
679 |
+
for iii in range(len(img_array)):
|
680 |
+
out.write(img_array[iii])
|
681 |
+
out.release()
|
682 |
+
# shutil.rmtree(os.path.join(cfg.experiment_dir, "video-png"))
|
683 |
+
|
684 |
+
print("The last loss is: {}".format(loss.item()))
|
685 |
+
|
686 |
|
687 |
if __name__ == "__main__":
|
688 |
|