Spaces:
Runtime error
Runtime error
""" | |
Here are some use cases: | |
python main.py --config config/all.yaml --experiment experiment_8x1 --signature demo1 --target data/demo1.png | |
""" | |
import pydiffvg | |
import torch | |
import cv2 | |
import matplotlib.pyplot as plt | |
import random | |
import argparse | |
import math | |
import errno | |
from tqdm import tqdm | |
from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR | |
from torch.nn.functional import adaptive_avg_pool2d | |
import warnings | |
warnings.filterwarnings("ignore") | |
import PIL | |
import PIL.Image | |
import os | |
import os.path as osp | |
import numpy as np | |
import numpy.random as npr | |
import shutil | |
import copy | |
# import skfmm | |
from xing_loss import xing_loss | |
import yaml | |
from easydict import EasyDict as edict | |
pydiffvg.set_print_timing(False) | |
gamma = 1.0 | |
########## | |
# helper # | |
########## | |
from utils import \ | |
get_experiment_id, \ | |
get_path_schedule, \ | |
edict_2_dict, \ | |
check_and_create_dir | |
def get_bezier_circle(radius=1, segments=4, bias=None): | |
points = [] | |
if bias is None: | |
bias = (random.random(), random.random()) | |
avg_degree = 360 / (segments*3) | |
for i in range(0, segments*3): | |
point = (np.cos(np.deg2rad(i * avg_degree)), | |
np.sin(np.deg2rad(i * avg_degree))) | |
points.append(point) | |
points = torch.tensor(points) | |
points = (points)*radius + torch.tensor(bias).unsqueeze(dim=0) | |
points = points.type(torch.FloatTensor) | |
return points | |
def get_sdf(phi, method='skfmm', **kwargs): | |
if method == 'skfmm': | |
import skfmm | |
phi = (phi-0.5)*2 | |
if (phi.max() <= 0) or (phi.min() >= 0): | |
return np.zeros(phi.shape).astype(np.float32) | |
sd = skfmm.distance(phi, dx=1) | |
flip_negative = kwargs.get('flip_negative', True) | |
if flip_negative: | |
sd = np.abs(sd) | |
truncate = kwargs.get('truncate', 10) | |
sd = np.clip(sd, -truncate, truncate) | |
# print(f"max sd value is: {sd.max()}") | |
zero2max = kwargs.get('zero2max', True) | |
if zero2max and flip_negative: | |
sd = sd.max() - sd | |
elif zero2max: | |
raise ValueError | |
normalize = kwargs.get('normalize', 'sum') | |
if normalize == 'sum': | |
sd /= sd.sum() | |
elif normalize == 'to1': | |
sd /= sd.max() | |
return sd | |
def parse_args(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument('--debug', action='store_true', default=False) | |
parser.add_argument("--config", type=str) | |
parser.add_argument("--experiment", type=str) | |
parser.add_argument("--seed", type=int) | |
parser.add_argument("--target", type=str, help="target image path") | |
parser.add_argument('--log_dir', metavar='DIR', default="log/debug") | |
parser.add_argument('--initial', type=str, default="random", choices=['random', 'circle']) | |
parser.add_argument('--signature', nargs='+', type=str) | |
parser.add_argument('--seginit', nargs='+', type=str) | |
parser.add_argument("--num_segments", type=int, default=4) | |
# parser.add_argument("--num_paths", type=str, default="1,1,1") | |
# parser.add_argument("--num_iter", type=int, default=500) | |
# parser.add_argument('--free', action='store_true') | |
# Please ensure that image resolution is divisible by pool_size; otherwise the performance would drop a lot. | |
# parser.add_argument('--pool_size', type=int, default=40, help="the pooled image size for next path initialization") | |
# parser.add_argument('--save_loss', action='store_true') | |
# parser.add_argument('--save_init', action='store_true') | |
# parser.add_argument('--save_image', action='store_true') | |
# parser.add_argument('--save_video', action='store_true') | |
# parser.add_argument('--print_weight', action='store_true') | |
# parser.add_argument('--circle_init_radius', type=float) | |
cfg = edict() | |
args = parser.parse_args() | |
cfg.debug = args.debug | |
cfg.config = args.config | |
cfg.experiment = args.experiment | |
cfg.seed = args.seed | |
cfg.target = args.target | |
cfg.log_dir = args.log_dir | |
cfg.initial = args.initial | |
cfg.signature = args.signature | |
# set cfg num_segments in command | |
cfg.num_segments = args.num_segments | |
if args.seginit is not None: | |
cfg.seginit = edict() | |
cfg.seginit.type = args.seginit[0] | |
if cfg.seginit.type == 'circle': | |
cfg.seginit.radius = float(args.seginit[1]) | |
return cfg | |
def ycrcb_conversion(im, format='[bs x 3 x 2D]', reverse=False): | |
mat = torch.FloatTensor([ | |
[ 65.481/255, 128.553/255, 24.966/255], # ranged_from [0, 219/255] | |
[-37.797/255, -74.203/255, 112.000/255], # ranged_from [-112/255, 112/255] | |
[112.000/255, -93.786/255, -18.214/255], # ranged_from [-112/255, 112/255] | |
]).to(im.device) | |
if reverse: | |
mat = mat.inverse() | |
if format == '[bs x 3 x 2D]': | |
im = im.permute(0, 2, 3, 1) | |
im = torch.matmul(im, mat.T) | |
im = im.permute(0, 3, 1, 2).contiguous() | |
return im | |
elif format == '[2D x 3]': | |
im = torch.matmul(im, mat.T) | |
return im | |
else: | |
raise ValueError | |
class random_coord_init(): | |
def __init__(self, canvas_size): | |
self.canvas_size = canvas_size | |
def __call__(self): | |
h, w = self.canvas_size | |
return [npr.uniform(0, 1)*w, npr.uniform(0, 1)*h] | |
class naive_coord_init(): | |
def __init__(self, pred, gt, format='[bs x c x 2D]', replace_sampling=True): | |
if isinstance(pred, torch.Tensor): | |
pred = pred.detach().cpu().numpy() | |
if isinstance(gt, torch.Tensor): | |
gt = gt.detach().cpu().numpy() | |
if format == '[bs x c x 2D]': | |
self.map = ((pred[0] - gt[0])**2).sum(0) | |
elif format == ['[2D x c]']: | |
self.map = ((pred - gt)**2).sum(-1) | |
else: | |
raise ValueError | |
self.replace_sampling = replace_sampling | |
def __call__(self): | |
coord = np.where(self.map == self.map.max()) | |
coord_h, coord_w = coord[0][0], coord[1][0] | |
if self.replace_sampling: | |
self.map[coord_h, coord_w] = -1 | |
return [coord_w, coord_h] | |
class sparse_coord_init(): | |
def __init__(self, pred, gt, format='[bs x c x 2D]', quantile_interval=200, nodiff_thres=0.1): | |
if isinstance(pred, torch.Tensor): | |
pred = pred.detach().cpu().numpy() | |
if isinstance(gt, torch.Tensor): | |
gt = gt.detach().cpu().numpy() | |
if format == '[bs x c x 2D]': | |
self.map = ((pred[0] - gt[0])**2).sum(0) | |
self.reference_gt = copy.deepcopy( | |
np.transpose(gt[0], (1, 2, 0))) | |
elif format == ['[2D x c]']: | |
self.map = (np.abs(pred - gt)).sum(-1) | |
self.reference_gt = copy.deepcopy(gt[0]) | |
else: | |
raise ValueError | |
# OptionA: Zero too small errors to avoid the error too small deadloop | |
self.map[self.map < nodiff_thres] = 0 | |
quantile_interval = np.linspace(0., 1., quantile_interval) | |
quantized_interval = np.quantile(self.map, quantile_interval) | |
# remove redundant | |
quantized_interval = np.unique(quantized_interval) | |
quantized_interval = sorted(quantized_interval[1:-1]) | |
self.map = np.digitize(self.map, quantized_interval, right=False) | |
self.map = np.clip(self.map, 0, 255).astype(np.uint8) | |
self.idcnt = {} | |
for idi in sorted(np.unique(self.map)): | |
self.idcnt[idi] = (self.map==idi).sum() | |
self.idcnt.pop(min(self.idcnt.keys())) | |
# remove smallest one to remove the correct region | |
def __call__(self): | |
if len(self.idcnt) == 0: | |
h, w = self.map.shape | |
return [npr.uniform(0, 1)*w, npr.uniform(0, 1)*h] | |
target_id = max(self.idcnt, key=self.idcnt.get) | |
_, component, cstats, ccenter = cv2.connectedComponentsWithStats( | |
(self.map==target_id).astype(np.uint8), connectivity=4) | |
# remove cid = 0, it is the invalid area | |
csize = [ci[-1] for ci in cstats[1:]] | |
target_cid = csize.index(max(csize))+1 | |
center = ccenter[target_cid][::-1] | |
coord = np.stack(np.where(component == target_cid)).T | |
dist = np.linalg.norm(coord-center, axis=1) | |
target_coord_id = np.argmin(dist) | |
coord_h, coord_w = coord[target_coord_id] | |
# replace_sampling | |
self.idcnt[target_id] -= max(csize) | |
if self.idcnt[target_id] == 0: | |
self.idcnt.pop(target_id) | |
self.map[component == target_cid] = 0 | |
return [coord_w, coord_h] | |
def init_shapes(num_paths, | |
num_segments, | |
canvas_size, | |
seginit_cfg, | |
shape_cnt, | |
pos_init_method=None, | |
trainable_stroke=False, | |
**kwargs): | |
shapes = [] | |
shape_groups = [] | |
h, w = canvas_size | |
# change path init location | |
if pos_init_method is None: | |
pos_init_method = random_coord_init(canvas_size=canvas_size) | |
for i in range(num_paths): | |
num_control_points = [2] * num_segments | |
if seginit_cfg.type=="random": | |
points = [] | |
p0 = pos_init_method() | |
color_ref = copy.deepcopy(p0) | |
points.append(p0) | |
for j in range(num_segments): | |
radius = seginit_cfg.radius | |
p1 = (p0[0] + radius * npr.uniform(-0.5, 0.5), | |
p0[1] + radius * npr.uniform(-0.5, 0.5)) | |
p2 = (p1[0] + radius * npr.uniform(-0.5, 0.5), | |
p1[1] + radius * npr.uniform(-0.5, 0.5)) | |
p3 = (p2[0] + radius * npr.uniform(-0.5, 0.5), | |
p2[1] + radius * npr.uniform(-0.5, 0.5)) | |
points.append(p1) | |
points.append(p2) | |
if j < num_segments - 1: | |
points.append(p3) | |
p0 = p3 | |
points = torch.FloatTensor(points) | |
# circle points initialization | |
elif seginit_cfg.type=="circle": | |
radius = seginit_cfg.radius | |
if radius is None: | |
radius = npr.uniform(0.5, 1) | |
center = pos_init_method() | |
color_ref = copy.deepcopy(center) | |
points = get_bezier_circle( | |
radius=radius, segments=num_segments, | |
bias=center) | |
path = pydiffvg.Path(num_control_points = torch.LongTensor(num_control_points), | |
points = points, | |
stroke_width = torch.tensor(0.0), | |
is_closed = True) | |
shapes.append(path) | |
# !!!!!!problem is here. the shape group shape_ids is wrong | |
if 'gt' in kwargs: | |
wref, href = color_ref | |
wref = max(0, min(int(wref), w-1)) | |
href = max(0, min(int(href), h-1)) | |
fill_color_init = list(gt[0, :, href, wref]) + [1.] | |
fill_color_init = torch.FloatTensor(fill_color_init) | |
stroke_color_init = torch.FloatTensor(npr.uniform(size=[4])) | |
else: | |
fill_color_init = torch.FloatTensor(npr.uniform(size=[4])) | |
stroke_color_init = torch.FloatTensor(npr.uniform(size=[4])) | |
path_group = pydiffvg.ShapeGroup( | |
shape_ids = torch.LongTensor([shape_cnt+i]), | |
fill_color = fill_color_init, | |
stroke_color = stroke_color_init, | |
) | |
shape_groups.append(path_group) | |
point_var = [] | |
color_var = [] | |
for path in shapes: | |
path.points.requires_grad = True | |
point_var.append(path.points) | |
for group in shape_groups: | |
group.fill_color.requires_grad = True | |
color_var.append(group.fill_color) | |
if trainable_stroke: | |
stroke_width_var = [] | |
stroke_color_var = [] | |
for path in shapes: | |
path.stroke_width.requires_grad = True | |
stroke_width_var.append(path.stroke_width) | |
for group in shape_groups: | |
group.stroke_color.requires_grad = True | |
stroke_color_var.append(group.stroke_color) | |
return shapes, shape_groups, point_var, color_var, stroke_width_var, stroke_color_var | |
else: | |
return shapes, shape_groups, point_var, color_var | |
class linear_decay_lrlambda_f(object): | |
def __init__(self, decay_every, decay_ratio): | |
self.decay_every = decay_every | |
self.decay_ratio = decay_ratio | |
def __call__(self, n): | |
decay_time = n//self.decay_every | |
decay_step = n %self.decay_every | |
lr_s = self.decay_ratio**decay_time | |
lr_e = self.decay_ratio**(decay_time+1) | |
r = decay_step/self.decay_every | |
lr = lr_s * (1-r) + lr_e * r | |
return lr | |
def main_func(target, experiment, cfg_arg): | |
with open(cfg_arg.config, 'r') as f: | |
cfg = yaml.load(f, Loader=yaml.FullLoader) | |
cfg_default = edict(cfg['default']) | |
cfg = edict(cfg[cfg_arg.experiment]) | |
cfg.update(cfg_default) | |
cfg.update(cfg_arg) | |
cfg.exid = get_experiment_id(cfg.debug) | |
cfg.experiment_dir = \ | |
osp.join(cfg.log_dir, '{}_{}'.format(cfg.exid, '_'.join(cfg.signature))) | |
cfg.target = target | |
cfg.experiment = experiment | |
configfile = osp.join(cfg.experiment_dir, 'config.yaml') | |
check_and_create_dir(configfile) | |
with open(osp.join(configfile), 'w') as f: | |
yaml.dump(edict_2_dict(cfg), f) | |
# Use GPU if available | |
pydiffvg.set_use_gpu(torch.cuda.is_available()) | |
device = pydiffvg.get_device() | |
gt = np.array(PIL.Image.open(cfg.target)) | |
print(f"Input image shape is: {gt.shape}") | |
if len(gt.shape) == 2: | |
print("Converting the gray-scale image to RGB.") | |
gt = gt.unsqueeze(dim=-1).repeat(1,1,3) | |
if gt.shape[2] == 4: | |
print("Input image includes alpha channel, simply dropout alpha channel.") | |
gt = gt[:, :, :3] | |
gt = (gt/255).astype(np.float32) | |
gt = torch.FloatTensor(gt).permute(2, 0, 1)[None].to(device) | |
if cfg.use_ycrcb: | |
gt = ycrcb_conversion(gt) | |
h, w = gt.shape[2:] | |
path_schedule = get_path_schedule(**cfg.path_schedule) | |
if cfg.seed is not None: | |
random.seed(cfg.seed) | |
npr.seed(cfg.seed) | |
torch.manual_seed(cfg.seed) | |
render = pydiffvg.RenderFunction.apply | |
shapes_record, shape_groups_record = [], [] | |
region_loss = None | |
loss_matrix = [] | |
para_point, para_color = {}, {} | |
if cfg.trainable.stroke: | |
para_stroke_width, para_stroke_color = {}, {} | |
pathn_record = [] | |
# Background | |
if cfg.trainable.bg: | |
# meancolor = gt.mean([2, 3])[0] | |
para_bg = torch.tensor([1., 1., 1.], requires_grad=True, device=device) | |
else: | |
if cfg.use_ycrcb: | |
para_bg = torch.tensor([219/255, 0, 0], requires_grad=False, device=device) | |
else: | |
para_bg = torch.tensor([1., 1., 1.], requires_grad=False, device=device) | |
################## | |
# start_training # | |
################## | |
loss_weight = None | |
loss_weight_keep = 0 | |
if cfg.coord_init.type == 'naive': | |
pos_init_method = naive_coord_init( | |
para_bg.view(1, -1, 1, 1).repeat(1, 1, h, w), gt) | |
elif cfg.coord_init.type == 'sparse': | |
pos_init_method = sparse_coord_init( | |
para_bg.view(1, -1, 1, 1).repeat(1, 1, h, w), gt) | |
elif cfg.coord_init.type == 'random': | |
pos_init_method = random_coord_init([h, w]) | |
else: | |
raise ValueError | |
lrlambda_f = linear_decay_lrlambda_f(cfg.num_iter, 0.4) | |
optim_schedular_dict = {} | |
for path_idx, pathn in enumerate(path_schedule): | |
loss_list = [] | |
print("=> Adding [{}] paths, [{}] ...".format(pathn, cfg.seginit.type)) | |
pathn_record.append(pathn) | |
pathn_record_str = '-'.join([str(i) for i in pathn_record]) | |
# initialize new shapes related stuffs. | |
if cfg.trainable.stroke: | |
shapes, shape_groups, point_var, color_var, stroke_width_var, stroke_color_var = init_shapes( | |
pathn, cfg.num_segments, (h, w), | |
cfg.seginit, len(shapes_record), | |
pos_init_method, | |
trainable_stroke=True, | |
gt=gt, ) | |
para_stroke_width[path_idx] = stroke_width_var | |
para_stroke_color[path_idx] = stroke_color_var | |
else: | |
shapes, shape_groups, point_var, color_var = init_shapes( | |
pathn, cfg.num_segments, (h, w), | |
cfg.seginit, len(shapes_record), | |
pos_init_method, | |
trainable_stroke=False, | |
gt=gt, ) | |
shapes_record += shapes | |
shape_groups_record += shape_groups | |
if cfg.save.init: | |
filename = os.path.join( | |
cfg.experiment_dir, "svg-init", | |
"{}-init.svg".format(pathn_record_str)) | |
check_and_create_dir(filename) | |
pydiffvg.save_svg( | |
filename, w, h, | |
shapes_record, shape_groups_record) | |
para = {} | |
if (cfg.trainable.bg) and (path_idx == 0): | |
para['bg'] = [para_bg] | |
para['point'] = point_var | |
para['color'] = color_var | |
if cfg.trainable.stroke: | |
para['stroke_width'] = stroke_width_var | |
para['stroke_color'] = stroke_color_var | |
pg = [{'params' : para[ki], 'lr' : cfg.lr_base[ki]} for ki in sorted(para.keys())] | |
optim = torch.optim.Adam(pg) | |
if cfg.trainable.record: | |
scheduler = LambdaLR( | |
optim, lr_lambda=lrlambda_f, last_epoch=-1) | |
else: | |
scheduler = LambdaLR( | |
optim, lr_lambda=lrlambda_f, last_epoch=cfg.num_iter) | |
optim_schedular_dict[path_idx] = (optim, scheduler) | |
# Inner loop training | |
t_range = tqdm(range(cfg.num_iter)) | |
for t in t_range: | |
for _, (optim, _) in optim_schedular_dict.items(): | |
optim.zero_grad() | |
# Forward pass: render the image. | |
scene_args = pydiffvg.RenderFunction.serialize_scene( | |
w, h, shapes_record, shape_groups_record) | |
img = render(w, h, 2, 2, t, None, *scene_args) | |
# Compose img with white background | |
img = img[:, :, 3:4] * img[:, :, :3] + \ | |
para_bg * (1 - img[:, :, 3:4]) | |
if cfg.save.video: | |
filename = os.path.join( | |
cfg.experiment_dir, "video-png", | |
"{}-iter{}.png".format(pathn_record_str, t)) | |
check_and_create_dir(filename) | |
if cfg.use_ycrcb: | |
imshow = ycrcb_conversion( | |
img, format='[2D x 3]', reverse=True).detach().cpu() | |
else: | |
imshow = img.detach().cpu() | |
pydiffvg.imwrite(imshow, filename, gamma=gamma) | |
x = img.unsqueeze(0).permute(0, 3, 1, 2) # HWC -> NCHW | |
if cfg.use_ycrcb: | |
color_reweight = torch.FloatTensor([255/219, 255/224, 255/255]).to(device) | |
loss = ((x-gt)*(color_reweight.view(1, -1, 1, 1)))**2 | |
else: | |
loss = ((x-gt)**2) | |
if cfg.loss.use_l1_loss: | |
loss = abs(x-gt) | |
if cfg.loss.use_distance_weighted_loss: | |
if cfg.use_ycrcb: | |
raise ValueError | |
shapes_forsdf = copy.deepcopy(shapes) | |
shape_groups_forsdf = copy.deepcopy(shape_groups) | |
for si in shapes_forsdf: | |
si.stroke_width = torch.FloatTensor([0]).to(device) | |
for sg_idx, sgi in enumerate(shape_groups_forsdf): | |
sgi.fill_color = torch.FloatTensor([1, 1, 1, 1]).to(device) | |
sgi.shape_ids = torch.LongTensor([sg_idx]).to(device) | |
sargs_forsdf = pydiffvg.RenderFunction.serialize_scene( | |
w, h, shapes_forsdf, shape_groups_forsdf) | |
with torch.no_grad(): | |
im_forsdf = render(w, h, 2, 2, 0, None, *sargs_forsdf) | |
# use alpha channel is a trick to get 0-1 image | |
im_forsdf = (im_forsdf[:, :, 3]).detach().cpu().numpy() | |
loss_weight = get_sdf(im_forsdf, normalize='to1') | |
loss_weight += loss_weight_keep | |
loss_weight = np.clip(loss_weight, 0, 1) | |
loss_weight = torch.FloatTensor(loss_weight).to(device) | |
if cfg.save.loss: | |
save_loss = loss.squeeze(dim=0).mean(dim=0,keepdim=False).cpu().detach().numpy() | |
save_weight = loss_weight.cpu().detach().numpy() | |
save_weighted_loss = save_loss*save_weight | |
# normalize to [0,1] | |
save_loss = (save_loss - np.min(save_loss))/np.ptp(save_loss) | |
save_weight = (save_weight - np.min(save_weight))/np.ptp(save_weight) | |
save_weighted_loss = (save_weighted_loss - np.min(save_weighted_loss))/np.ptp(save_weighted_loss) | |
# save | |
plt.imshow(save_loss, cmap='Reds') | |
plt.axis('off') | |
# plt.colorbar() | |
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-mseloss.png".format(pathn_record_str, t)) | |
check_and_create_dir(filename) | |
plt.savefig(filename, dpi=800) | |
plt.close() | |
plt.imshow(save_weight, cmap='Greys') | |
plt.axis('off') | |
# plt.colorbar() | |
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-sdfweight.png".format(pathn_record_str, t)) | |
plt.savefig(filename, dpi=800) | |
plt.close() | |
plt.imshow(save_weighted_loss, cmap='Reds') | |
plt.axis('off') | |
# plt.colorbar() | |
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-weightedloss.png".format(pathn_record_str, t)) | |
plt.savefig(filename, dpi=800) | |
plt.close() | |
if loss_weight is None: | |
loss = loss.sum(1).mean() | |
else: | |
loss = (loss.sum(1)*loss_weight).mean() | |
# if (cfg.loss.bis_loss_weight is not None) and (cfg.loss.bis_loss_weight > 0): | |
# loss_bis = bezier_intersection_loss(point_var[0]) * cfg.loss.bis_loss_weight | |
# loss = loss + loss_bis | |
if (cfg.loss.xing_loss_weight is not None) \ | |
and (cfg.loss.xing_loss_weight > 0): | |
loss_xing = xing_loss(point_var) * cfg.loss.xing_loss_weight | |
loss = loss + loss_xing | |
loss_list.append(loss.item()) | |
t_range.set_postfix({'loss': loss.item()}) | |
loss.backward() | |
# step | |
for _, (optim, scheduler) in optim_schedular_dict.items(): | |
optim.step() | |
scheduler.step() | |
for group in shape_groups_record: | |
group.fill_color.data.clamp_(0.0, 1.0) | |
if cfg.loss.use_distance_weighted_loss: | |
loss_weight_keep = loss_weight.detach().cpu().numpy() * 1 | |
if not cfg.trainable.record: | |
for _, pi in pg.items(): | |
for ppi in pi: | |
pi.require_grad = False | |
optim_schedular_dict = {} | |
if cfg.save.image: | |
filename = os.path.join( | |
cfg.experiment_dir, "demo-png", "{}.png".format(pathn_record_str)) | |
check_and_create_dir(filename) | |
if cfg.use_ycrcb: | |
imshow = ycrcb_conversion( | |
img, format='[2D x 3]', reverse=True).detach().cpu() | |
else: | |
imshow = img.detach().cpu() | |
pydiffvg.imwrite(imshow, filename, gamma=gamma) | |
if cfg.save.output: | |
filename = os.path.join( | |
cfg.experiment_dir, "output-svg", "{}.svg".format(pathn_record_str)) | |
check_and_create_dir(filename) | |
pydiffvg.save_svg(filename, w, h, shapes_record, shape_groups_record) | |
loss_matrix.append(loss_list) | |
# calculate the pixel loss | |
# pixel_loss = ((x-gt)**2).sum(dim=1, keepdim=True).sqrt_() # [N,1,H, W] | |
# region_loss = adaptive_avg_pool2d(pixel_loss, cfg.region_loss_pool_size) | |
# loss_weight = torch.softmax(region_loss.reshape(1, 1, -1), dim=-1)\ | |
# .reshape_as(region_loss) | |
pos_init_method = naive_coord_init(x, gt) | |
if cfg.coord_init.type == 'naive': | |
pos_init_method = naive_coord_init(x, gt) | |
elif cfg.coord_init.type == 'sparse': | |
pos_init_method = sparse_coord_init(x, gt) | |
elif cfg.coord_init.type == 'random': | |
pos_init_method = random_coord_init([h, w]) | |
else: | |
raise ValueError | |
if cfg.save.video: | |
print("saving iteration video...") | |
img_array = [] | |
for ii in range(0, cfg.num_iter): | |
filename = os.path.join( | |
cfg.experiment_dir, "video-png", | |
"{}-iter{}.png".format(pathn_record_str, ii)) | |
img = cv2.imread(filename) | |
# cv2.putText( | |
# img, "Path:{} \nIteration:{}".format(pathn_record_str, ii), | |
# (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) | |
img_array.append(img) | |
videoname = os.path.join( | |
cfg.experiment_dir, "video-avi", | |
"{}.avi".format(pathn_record_str)) | |
check_and_create_dir(videoname) | |
out = cv2.VideoWriter( | |
videoname, | |
# cv2.VideoWriter_fourcc(*'mp4v'), | |
cv2.VideoWriter_fourcc(*'FFV1'), | |
20.0, (w, h)) | |
for iii in range(len(img_array)): | |
out.write(img_array[iii]) | |
out.release() | |
# shutil.rmtree(os.path.join(cfg.experiment_dir, "video-png")) | |
print("The last loss is: {}".format(loss.item())) | |
if __name__ == "__main__": | |
############### | |
# make config # | |
############### | |
cfg_arg = parse_args() | |
with open(cfg_arg.config, 'r') as f: | |
cfg = yaml.load(f, Loader=yaml.FullLoader) | |
cfg_default = edict(cfg['default']) | |
cfg = edict(cfg[cfg_arg.experiment]) | |
cfg.update(cfg_default) | |
cfg.update(cfg_arg) | |
cfg.exid = get_experiment_id(cfg.debug) | |
cfg.experiment_dir = \ | |
osp.join(cfg.log_dir, '{}_{}'.format(cfg.exid, '_'.join(cfg.signature))) | |
configfile = osp.join(cfg.experiment_dir, 'config.yaml') | |
check_and_create_dir(configfile) | |
with open(osp.join(configfile), 'w') as f: | |
yaml.dump(edict_2_dict(cfg), f) | |
# Use GPU if available | |
pydiffvg.set_use_gpu(torch.cuda.is_available()) | |
device = pydiffvg.get_device() | |
gt = np.array(PIL.Image.open(cfg.target)) | |
print(f"Input image shape is: {gt.shape}") | |
if len(gt.shape) == 2: | |
print("Converting the gray-scale image to RGB.") | |
gt = gt.unsqueeze(dim=-1).repeat(1,1,3) | |
if gt.shape[2] == 4: | |
print("Input image includes alpha channel, simply dropout alpha channel.") | |
gt = gt[:, :, :3] | |
gt = (gt/255).astype(np.float32) | |
gt = torch.FloatTensor(gt).permute(2, 0, 1)[None].to(device) | |
if cfg.use_ycrcb: | |
gt = ycrcb_conversion(gt) | |
h, w = gt.shape[2:] | |
path_schedule = get_path_schedule(**cfg.path_schedule) | |
if cfg.seed is not None: | |
random.seed(cfg.seed) | |
npr.seed(cfg.seed) | |
torch.manual_seed(cfg.seed) | |
render = pydiffvg.RenderFunction.apply | |
shapes_record, shape_groups_record = [], [] | |
region_loss = None | |
loss_matrix = [] | |
para_point, para_color = {}, {} | |
if cfg.trainable.stroke: | |
para_stroke_width, para_stroke_color = {}, {} | |
pathn_record = [] | |
# Background | |
if cfg.trainable.bg: | |
# meancolor = gt.mean([2, 3])[0] | |
para_bg = torch.tensor([1., 1., 1.], requires_grad=True, device=device) | |
else: | |
if cfg.use_ycrcb: | |
para_bg = torch.tensor([219/255, 0, 0], requires_grad=False, device=device) | |
else: | |
para_bg = torch.tensor([1., 1., 1.], requires_grad=False, device=device) | |
################## | |
# start_training # | |
################## | |
loss_weight = None | |
loss_weight_keep = 0 | |
if cfg.coord_init.type == 'naive': | |
pos_init_method = naive_coord_init( | |
para_bg.view(1, -1, 1, 1).repeat(1, 1, h, w), gt) | |
elif cfg.coord_init.type == 'sparse': | |
pos_init_method = sparse_coord_init( | |
para_bg.view(1, -1, 1, 1).repeat(1, 1, h, w), gt) | |
elif cfg.coord_init.type == 'random': | |
pos_init_method = random_coord_init([h, w]) | |
else: | |
raise ValueError | |
lrlambda_f = linear_decay_lrlambda_f(cfg.num_iter, 0.4) | |
optim_schedular_dict = {} | |
for path_idx, pathn in enumerate(path_schedule): | |
loss_list = [] | |
print("=> Adding [{}] paths, [{}] ...".format(pathn, cfg.seginit.type)) | |
pathn_record.append(pathn) | |
pathn_record_str = '-'.join([str(i) for i in pathn_record]) | |
# initialize new shapes related stuffs. | |
if cfg.trainable.stroke: | |
shapes, shape_groups, point_var, color_var, stroke_width_var, stroke_color_var = init_shapes( | |
pathn, cfg.num_segments, (h, w), | |
cfg.seginit, len(shapes_record), | |
pos_init_method, | |
trainable_stroke=True, | |
gt=gt, ) | |
para_stroke_width[path_idx] = stroke_width_var | |
para_stroke_color[path_idx] = stroke_color_var | |
else: | |
shapes, shape_groups, point_var, color_var = init_shapes( | |
pathn, cfg.num_segments, (h, w), | |
cfg.seginit, len(shapes_record), | |
pos_init_method, | |
trainable_stroke=False, | |
gt=gt, ) | |
shapes_record += shapes | |
shape_groups_record += shape_groups | |
if cfg.save.init: | |
filename = os.path.join( | |
cfg.experiment_dir, "svg-init", | |
"{}-init.svg".format(pathn_record_str)) | |
check_and_create_dir(filename) | |
pydiffvg.save_svg( | |
filename, w, h, | |
shapes_record, shape_groups_record) | |
para = {} | |
if (cfg.trainable.bg) and (path_idx == 0): | |
para['bg'] = [para_bg] | |
para['point'] = point_var | |
para['color'] = color_var | |
if cfg.trainable.stroke: | |
para['stroke_width'] = stroke_width_var | |
para['stroke_color'] = stroke_color_var | |
pg = [{'params' : para[ki], 'lr' : cfg.lr_base[ki]} for ki in sorted(para.keys())] | |
optim = torch.optim.Adam(pg) | |
if cfg.trainable.record: | |
scheduler = LambdaLR( | |
optim, lr_lambda=lrlambda_f, last_epoch=-1) | |
else: | |
scheduler = LambdaLR( | |
optim, lr_lambda=lrlambda_f, last_epoch=cfg.num_iter) | |
optim_schedular_dict[path_idx] = (optim, scheduler) | |
# Inner loop training | |
t_range = tqdm(range(cfg.num_iter)) | |
for t in t_range: | |
for _, (optim, _) in optim_schedular_dict.items(): | |
optim.zero_grad() | |
# Forward pass: render the image. | |
scene_args = pydiffvg.RenderFunction.serialize_scene( | |
w, h, shapes_record, shape_groups_record) | |
img = render(w, h, 2, 2, t, None, *scene_args) | |
# Compose img with white background | |
img = img[:, :, 3:4] * img[:, :, :3] + \ | |
para_bg * (1 - img[:, :, 3:4]) | |
if cfg.save.video: | |
filename = os.path.join( | |
cfg.experiment_dir, "video-png", | |
"{}-iter{}.png".format(pathn_record_str, t)) | |
check_and_create_dir(filename) | |
if cfg.use_ycrcb: | |
imshow = ycrcb_conversion( | |
img, format='[2D x 3]', reverse=True).detach().cpu() | |
else: | |
imshow = img.detach().cpu() | |
pydiffvg.imwrite(imshow, filename, gamma=gamma) | |
x = img.unsqueeze(0).permute(0, 3, 1, 2) # HWC -> NCHW | |
if cfg.use_ycrcb: | |
color_reweight = torch.FloatTensor([255/219, 255/224, 255/255]).to(device) | |
loss = ((x-gt)*(color_reweight.view(1, -1, 1, 1)))**2 | |
else: | |
loss = ((x-gt)**2) | |
if cfg.loss.use_l1_loss: | |
loss = abs(x-gt) | |
if cfg.loss.use_distance_weighted_loss: | |
if cfg.use_ycrcb: | |
raise ValueError | |
shapes_forsdf = copy.deepcopy(shapes) | |
shape_groups_forsdf = copy.deepcopy(shape_groups) | |
for si in shapes_forsdf: | |
si.stroke_width = torch.FloatTensor([0]).to(device) | |
for sg_idx, sgi in enumerate(shape_groups_forsdf): | |
sgi.fill_color = torch.FloatTensor([1, 1, 1, 1]).to(device) | |
sgi.shape_ids = torch.LongTensor([sg_idx]).to(device) | |
sargs_forsdf = pydiffvg.RenderFunction.serialize_scene( | |
w, h, shapes_forsdf, shape_groups_forsdf) | |
with torch.no_grad(): | |
im_forsdf = render(w, h, 2, 2, 0, None, *sargs_forsdf) | |
# use alpha channel is a trick to get 0-1 image | |
im_forsdf = (im_forsdf[:, :, 3]).detach().cpu().numpy() | |
loss_weight = get_sdf(im_forsdf, normalize='to1') | |
loss_weight += loss_weight_keep | |
loss_weight = np.clip(loss_weight, 0, 1) | |
loss_weight = torch.FloatTensor(loss_weight).to(device) | |
if cfg.save.loss: | |
save_loss = loss.squeeze(dim=0).mean(dim=0,keepdim=False).cpu().detach().numpy() | |
save_weight = loss_weight.cpu().detach().numpy() | |
save_weighted_loss = save_loss*save_weight | |
# normalize to [0,1] | |
save_loss = (save_loss - np.min(save_loss))/np.ptp(save_loss) | |
save_weight = (save_weight - np.min(save_weight))/np.ptp(save_weight) | |
save_weighted_loss = (save_weighted_loss - np.min(save_weighted_loss))/np.ptp(save_weighted_loss) | |
# save | |
plt.imshow(save_loss, cmap='Reds') | |
plt.axis('off') | |
# plt.colorbar() | |
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-mseloss.png".format(pathn_record_str, t)) | |
check_and_create_dir(filename) | |
plt.savefig(filename, dpi=800) | |
plt.close() | |
plt.imshow(save_weight, cmap='Greys') | |
plt.axis('off') | |
# plt.colorbar() | |
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-sdfweight.png".format(pathn_record_str, t)) | |
plt.savefig(filename, dpi=800) | |
plt.close() | |
plt.imshow(save_weighted_loss, cmap='Reds') | |
plt.axis('off') | |
# plt.colorbar() | |
filename = os.path.join(cfg.experiment_dir, "loss", "{}-iter{}-weightedloss.png".format(pathn_record_str, t)) | |
plt.savefig(filename, dpi=800) | |
plt.close() | |
if loss_weight is None: | |
loss = loss.sum(1).mean() | |
else: | |
loss = (loss.sum(1)*loss_weight).mean() | |
# if (cfg.loss.bis_loss_weight is not None) and (cfg.loss.bis_loss_weight > 0): | |
# loss_bis = bezier_intersection_loss(point_var[0]) * cfg.loss.bis_loss_weight | |
# loss = loss + loss_bis | |
if (cfg.loss.xing_loss_weight is not None) \ | |
and (cfg.loss.xing_loss_weight > 0): | |
loss_xing = xing_loss(point_var) * cfg.loss.xing_loss_weight | |
loss = loss + loss_xing | |
loss_list.append(loss.item()) | |
t_range.set_postfix({'loss': loss.item()}) | |
loss.backward() | |
# step | |
for _, (optim, scheduler) in optim_schedular_dict.items(): | |
optim.step() | |
scheduler.step() | |
for group in shape_groups_record: | |
group.fill_color.data.clamp_(0.0, 1.0) | |
if cfg.loss.use_distance_weighted_loss: | |
loss_weight_keep = loss_weight.detach().cpu().numpy() * 1 | |
if not cfg.trainable.record: | |
for _, pi in pg.items(): | |
for ppi in pi: | |
pi.require_grad = False | |
optim_schedular_dict = {} | |
if cfg.save.image: | |
filename = os.path.join( | |
cfg.experiment_dir, "demo-png", "{}.png".format(pathn_record_str)) | |
check_and_create_dir(filename) | |
if cfg.use_ycrcb: | |
imshow = ycrcb_conversion( | |
img, format='[2D x 3]', reverse=True).detach().cpu() | |
else: | |
imshow = img.detach().cpu() | |
pydiffvg.imwrite(imshow, filename, gamma=gamma) | |
if cfg.save.output: | |
filename = os.path.join( | |
cfg.experiment_dir, "output-svg", "{}.svg".format(pathn_record_str)) | |
check_and_create_dir(filename) | |
pydiffvg.save_svg(filename, w, h, shapes_record, shape_groups_record) | |
loss_matrix.append(loss_list) | |
# calculate the pixel loss | |
# pixel_loss = ((x-gt)**2).sum(dim=1, keepdim=True).sqrt_() # [N,1,H, W] | |
# region_loss = adaptive_avg_pool2d(pixel_loss, cfg.region_loss_pool_size) | |
# loss_weight = torch.softmax(region_loss.reshape(1, 1, -1), dim=-1)\ | |
# .reshape_as(region_loss) | |
pos_init_method = naive_coord_init(x, gt) | |
if cfg.coord_init.type == 'naive': | |
pos_init_method = naive_coord_init(x, gt) | |
elif cfg.coord_init.type == 'sparse': | |
pos_init_method = sparse_coord_init(x, gt) | |
elif cfg.coord_init.type == 'random': | |
pos_init_method = random_coord_init([h, w]) | |
else: | |
raise ValueError | |
if cfg.save.video: | |
print("saving iteration video...") | |
img_array = [] | |
for ii in range(0, cfg.num_iter): | |
filename = os.path.join( | |
cfg.experiment_dir, "video-png", | |
"{}-iter{}.png".format(pathn_record_str, ii)) | |
img = cv2.imread(filename) | |
# cv2.putText( | |
# img, "Path:{} \nIteration:{}".format(pathn_record_str, ii), | |
# (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) | |
img_array.append(img) | |
videoname = os.path.join( | |
cfg.experiment_dir, "video-avi", | |
"{}.avi".format(pathn_record_str)) | |
check_and_create_dir(videoname) | |
out = cv2.VideoWriter( | |
videoname, | |
# cv2.VideoWriter_fourcc(*'mp4v'), | |
cv2.VideoWriter_fourcc(*'FFV1'), | |
20.0, (w, h)) | |
for iii in range(len(img_array)): | |
out.write(img_array[iii]) | |
out.release() | |
# shutil.rmtree(os.path.join(cfg.experiment_dir, "video-png")) | |
print("The last loss is: {}".format(loss.item())) | |