def save(self, x, save_dir, name):
z, log_var = self.encode(x)
all_points = self.decode(z)
# print(all_points.std(dim=1))
# all_points = ((all_points-0.5)*2 + 0.5)*self.imsize
# if type(self.sort_idx) == type(None):
# angles = torch.atan(all_points[:,:,1]/all_points[:,:,0]).detach()
# self.sort_idx = torch.argsort(angles, dim=1)
# Process the batch sequentially
outputs = []
for k in range(1):
# Get point parameters from network
shapes = []
shape_groups = []
points = all_points[k].cpu()#[self.sort_idx[k]]
color = torch.cat([torch.tensor([0,0,0,1]),])
num_ctrl_pts = torch.zeros(self.curves, dtype=torch.int32) + 2
path = pydiffvg.Path(
num_control_points=num_ctrl_pts, points=points,
is_closed=True)
shapes.append(path)
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(shapes) - 1]),
fill_color=color,
stroke_color=color)
shape_groups.append(path_group)
pydiffvg.save_svg(f"{save_dir}{name}/{name}.svg",
self.imsize, self.imsize, shapes, shape_groups)
def run(n_paths_original, im_size_original, n_steps, layer_name, layer_index):
upscale_x = 4
upscale_y = 3
name = "result_" + "_".join([
f"{k}{v}" for k, v in zip(
["n_paths", "im_size", "n_steps", "layer_name", "layer_index"], [
n_paths_original, im_size_original, n_steps, layer_name,
layer_index
])
])
with SummaryWriter(log_dir=f"./logs/{name}", comment=name) as writer:
gen = Generator(n_paths_original, im_size_original, im_size_original)
optimizer = CurveOptimizer(n_steps,
im_size_original,
im_size_original,
gen.gen_func(),
gen_vgg16_excitation_func(
layer_name, layer_index),
scale=(0.9, 1.05),
n_augms=8)
shapes, shape_groups = optimizer.gen_and_optimize(
writer, color_optimisation_activated=False)
filename = "./" + name.replace(".", "_") + ".svg"
pydiffvg.save_svg(filename, im_size_original, im_size_original, shapes,
shape_groups)
n_paths = upscale_x * upscale_y * n_paths_original
im_size_x = upscale_x * im_size_original
im_size_y = upscale_y * im_size_original
n_steps *= 2
large_name = "result" + "_".join([
f"{k}{v}" for k, v in
zip(["n_paths", "im_size", "n_steps", "layer_name", "layer_index"],
[n_paths, upscale_x, n_steps, layer_name, layer_index])
])
with SummaryWriter(log_dir=f"./logs/{large_name}",
comment=large_name) as writer:
gen = ScaledSvgGen(filename, upscale_y, upscale_x)
optimizer = CurveOptimizer(
n_steps,
im_size_y,
im_size_x,
gen.gen_func(),
gen_vgg16_excitation_func(layer_name, layer_index),
scale=[0.8 * 1. / upscale_y, 1.2 * 1. / upscale_x],
n_augms=12)
shapes, shape_groups = optimizer.gen_and_optimize(
writer, color_optimisation_activated=False)
filename_large = "./" + large_name.replace(".", "_") + ".svg"
pydiffvg.save_svg(filename_large, im_size_y, im_size_x, shapes,
shape_groups)
def run(n_paths, im_size, n_steps, img_path):
gen = Generator(n_paths, im_size, im_size)
optimizer = CurveOptimizer(n_steps, im_size, im_size, gen.gen_func(),
gen_vgg16_mimick(img_path))
shapes, shape_groups = optimizer.gen_and_optimize()
name = "result_" + "_".join([
f"{k}{v}" for k, v in zip(["n_paths", "im_size", "n_steps", "file"], [
n_paths, im_size, n_steps,
os.path.splitext(os.path.basename(img_path))[0]
])
]) + ".svg"
pydiffvg.save_svg(name, im_size, im_size, shapes, shape_groups)
def run(n_paths, im_size, n_steps, prompt):
name = "result_" + "_".join([
f"{k}{v}" for k, v in zip(["n_paths", "im_size", "n_steps", "prompt"],
[n_paths, im_size, n_steps, prompt])
])
with SummaryWriter(log_dir=f"./logs/{name}", comment=name) as writer:
gen = Generator(n_paths, im_size, im_size, allow_color=True)
fn = ClipForwardFunc(model, 4, "a boat on the sea").gen_func()
optimizer = CurveOptimizer(n_steps, im_size, im_size, gen.gen_func(),
fn)
shapes, shape_groups = optimizer.gen_and_optimize(
writer, color_optimisation_activated=True)
pydiffvg.save_svg(name + ".svg", im_size, im_size, shapes,
shape_groups)
def sample(gen,
encoder_A,
encoder_B,
train_x_loader,
train_y_loader,
resize=128,
svg=False):
gen.eval()
encoder_A.eval()
encoder_B.eval()
gen.imsize = resize
for batch, (batch_x, batch_y) in tqdm(
enumerate(zip(train_x_loader, train_y_loader), start=1)):
source = batch_x['source'].to(device)
target = batch_x['target'].to(device)
real = batch_y['target'].to(device)
# Generate encoding
fixed_z = encoder_A(real.unsqueeze(2))
fixed_z = encoder_B(fixed_z.squeeze(2))
z_dim_size = fixed_z.shape[1] * fixed_z.shape[2] * fixed_z.shape[3]
fixed_z = fixed_z.view(-1, 1, z_dim_size)
z = fixed_z.repeat(1, 52, 1) # shape = (bs*52, z_dim)
z = z.view(-1, z_dim_size)
glyph_one_hot = torch.eye(52).repeat(fixed_z.shape[0], 1).to(
device) # shape = (52*bs, 52)
z = torch.cat([z, glyph_one_hot], dim=1)
# Decode
if svg:
scenes = gen.get_vector(z)
for i, scene in enumerate(scenes):
pydiffvg.save_svg(f'svg/{i}.svg', *scene)
if i == 30:
break
# else:
gen_output_t = gen(z) # shape = (52*bs, resize, resize)
gen_output_t = gen_output_t.view(-1, 52, resize, resize)
for i in range(len(source)):
util.save_image_grid(f'svg/real_{i}.jpg',
real[i, :, :, :].detach().cpu().numpy() * 255)
util.save_image_grid(
f'svg/fake_{i}.jpg',
gen_output_t[i, :, :, :].detach().cpu().numpy() * 255)
break
def run(n_paths, im_size, n_steps, layer_name, layer_index):
name = "result_" + "_".join([
f"{k}{v}" for k, v in
zip(["n_paths", "im_size", "n_steps", "layer_name", "layer_index"],
[n_paths, im_size, n_steps, layer_name, layer_index])
])
with SummaryWriter(log_dir=f"./logs/{name}", comment=name) as writer:
gen = Generator(n_paths, im_size, im_size)
optimizer = CurveOptimizer(
n_steps, im_size, im_size, gen.gen_func(),
gen_vgg16_excitation_func(layer_name, layer_index))
shapes, shape_groups = optimizer.gen_and_optimize(
writer, color_optimisation_activated=False)
pydiffvg.save_svg(name + ".svg", im_size, im_size, shapes,
shape_groups)
def save_scene(scn, path):
os.makedirs(os.path.dirname(path), exist_ok=True)
pydiffvg.save_svg(path, *scn, use_gamma=False)
def main(args):
# Use GPU if available
pydiffvg.set_use_gpu(torch.cuda.is_available())
perception_loss = ttools.modules.LPIPS().to(pydiffvg.get_device())
#target = torch.from_numpy(skimage.io.imread('imgs/lena.png')).to(torch.float32) / 255.0
target = torch.from_numpy(skimage.io.imread(args.target)).to(
torch.float32) / 255.0
target = target.pow(gamma)
target = target.to(pydiffvg.get_device())
target = target.unsqueeze(0)
target = target.permute(0, 3, 1, 2) # NHWC -> NCHW
#target = torch.nn.functional.interpolate(target, size = [256, 256], mode = 'area')
canvas_width, canvas_height = target.shape[3], target.shape[2]
num_paths = args.num_paths
max_width = args.max_width
random.seed(1234)
torch.manual_seed(1234)
shapes = []
shape_groups = []
if args.use_blob:
for i in range(num_paths):
num_segments = random.randint(3, 5)
num_control_points = torch.zeros(num_segments,
dtype=torch.int32) + 2
points = []
p0 = (random.random(), random.random())
points.append(p0)
for j in range(num_segments):
radius = 0.05
p1 = (p0[0] + radius * (random.random() - 0.5),
p0[1] + radius * (random.random() - 0.5))
p2 = (p1[0] + radius * (random.random() - 0.5),
p1[1] + radius * (random.random() - 0.5))
p3 = (p2[0] + radius * (random.random() - 0.5),
p2[1] + radius * (random.random() - 0.5))
points.append(p1)
points.append(p2)
if j < num_segments - 1:
points.append(p3)
p0 = p3
points = torch.tensor(points)
points[:, 0] *= canvas_width
points[:, 1] *= canvas_height
path = pydiffvg.Path(num_control_points=num_control_points,
points=points,
stroke_width=torch.tensor(1.0),
is_closed=True)
shapes.append(path)
path_group = pydiffvg.ShapeGroup(shape_ids=torch.tensor(
[len(shapes) - 1]),
fill_color=torch.tensor([
random.random(),
random.random(),
random.random(),
random.random()
]))
shape_groups.append(path_group)
else:
for i in range(num_paths):
num_segments = random.randint(1, 3)
num_control_points = torch.zeros(num_segments,
dtype=torch.int32) + 2
points = []
p0 = (random.random(), random.random())
points.append(p0)
for j in range(num_segments):
radius = 0.05
p1 = (p0[0] + radius * (random.random() - 0.5),
p0[1] + radius * (random.random() - 0.5))
p2 = (p1[0] + radius * (random.random() - 0.5),
p1[1] + radius * (random.random() - 0.5))
p3 = (p2[0] + radius * (random.random() - 0.5),
p2[1] + radius * (random.random() - 0.5))
points.append(p1)
points.append(p2)
points.append(p3)
p0 = p3
points = torch.tensor(points)
points[:, 0] *= canvas_width
points[:, 1] *= canvas_height
#points = torch.rand(3 * num_segments + 1, 2) * min(canvas_width, canvas_height)
path = pydiffvg.Path(num_control_points=num_control_points,
points=points,
stroke_width=torch.tensor(1.0),
is_closed=False)
shapes.append(path)
path_group = pydiffvg.ShapeGroup(shape_ids=torch.tensor(
[len(shapes) - 1]),
fill_color=None,
stroke_color=torch.tensor([
random.random(),
random.random(),
random.random(),
random.random()
]))
shape_groups.append(path_group)
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
render = pydiffvg.RenderFunction.apply
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None,
*scene_args)
pydiffvg.imwrite(img.cpu(),
'results/painterly_rendering/init.png',
gamma=gamma)
points_vars = []
stroke_width_vars = []
color_vars = []
for path in shapes:
path.points.requires_grad = True
points_vars.append(path.points)
if not args.use_blob:
for path in shapes:
path.stroke_width.requires_grad = True
stroke_width_vars.append(path.stroke_width)
if args.use_blob:
for group in shape_groups:
group.fill_color.requires_grad = True
color_vars.append(group.fill_color)
else:
for group in shape_groups:
group.stroke_color.requires_grad = True
color_vars.append(group.stroke_color)
# Optimize
points_optim = torch.optim.Adam(points_vars, lr=1.0)
if len(stroke_width_vars) > 0:
width_optim = torch.optim.Adam(stroke_width_vars, lr=0.1)
color_optim = torch.optim.Adam(color_vars, lr=0.01)
# Adam iterations.
for t in range(args.num_iter):
print('iteration:', t)
points_optim.zero_grad()
if len(stroke_width_vars) > 0:
width_optim.zero_grad()
color_optim.zero_grad()
# Forward pass: render the image.
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
t, # seed
None,
*scene_args)
# Compose img with white background
img = img[:, :, 3:4] * img[:, :, :3] + torch.ones(
img.shape[0], img.shape[1], 3,
device=pydiffvg.get_device()) * (1 - img[:, :, 3:4])
# Save the intermediate render.
pydiffvg.imwrite(img.cpu(),
'results/painterly_rendering/iter_{}.png'.format(t),
gamma=gamma)
img = img[:, :, :3]
# Convert img from HWC to NCHW
img = img.unsqueeze(0)
img = img.permute(0, 3, 1, 2) # NHWC -> NCHW
if args.use_lpips_loss:
loss = perception_loss(
img, target) + (img.mean() - target.mean()).pow(2)
else:
loss = (img - target).pow(2).mean()
print('render loss:', loss.item())
# Backpropagate the gradients.
loss.backward()
# Take a gradient descent step.
points_optim.step()
if len(stroke_width_vars) > 0:
width_optim.step()
color_optim.step()
if len(stroke_width_vars) > 0:
for path in shapes:
path.stroke_width.data.clamp_(1.0, max_width)
if args.use_blob:
for group in shape_groups:
group.fill_color.data.clamp_(0.0, 1.0)
else:
for group in shape_groups:
group.stroke_color.data.clamp_(0.0, 1.0)
if t % 10 == 0 or t == args.num_iter - 1:
pydiffvg.save_svg(
'results/painterly_rendering/iter_{}.svg'.format(t),
canvas_width, canvas_height, shapes, shape_groups)
# Render the final result.
img = render(
target.shape[1], # width
target.shape[0], # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None,
*scene_args)
# Save the intermediate render.
pydiffvg.imwrite(img.cpu(),
'results/painterly_rendering/final.png'.format(t),
gamma=gamma)
# Convert the intermediate renderings to a video.
from subprocess import call
call([
"ffmpeg", "-framerate", "24", "-i",
"results/painterly_rendering/iter_%d.png", "-vb", "20M",
"results/painterly_rendering/out.mp4"
])
def main(args):
perception_loss = ttools.modules.LPIPS().to(pydiffvg.get_device())
target = torch.from_numpy(skimage.io.imread(args.target)).to(
torch.float32) / 255.0
target = target.pow(gamma)
target = target.to(pydiffvg.get_device())
target = target.unsqueeze(0)
target = target.permute(0, 3, 1, 2) # NHWC -> NCHW
canvas_width, canvas_height, shapes, shape_groups = \
pydiffvg.svg_to_scene(args.svg)
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
render = pydiffvg.RenderFunction.apply
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None, # bg
*scene_args)
# The output image is in linear RGB space. Do Gamma correction before saving the image.
pydiffvg.imwrite(img.cpu(), 'results/refine_svg/init.png', gamma=gamma)
points_vars = []
for path in shapes:
path.points.requires_grad = True
points_vars.append(path.points)
color_vars = {}
for group in shape_groups:
group.fill_color.requires_grad = True
color_vars[group.fill_color.data_ptr()] = group.fill_color
color_vars = list(color_vars.values())
# Optimize
points_optim = torch.optim.Adam(points_vars, lr=1.0)
color_optim = torch.optim.Adam(color_vars, lr=0.01)
# Adam iterations.
for t in range(args.num_iter):
print('iteration:', t)
points_optim.zero_grad()
color_optim.zero_grad()
# Forward pass: render the image.
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None, # bg
*scene_args)
# Compose img with white background
img = img[:, :, 3:4] * img[:, :, :3] + torch.ones(
img.shape[0], img.shape[1], 3,
device=pydiffvg.get_device()) * (1 - img[:, :, 3:4])
# Save the intermediate render.
pydiffvg.imwrite(img.cpu(),
'results/refine_svg/iter_{}.png'.format(t),
gamma=gamma)
img = img[:, :, :3]
# Convert img from HWC to NCHW
img = img.unsqueeze(0)
img = img.permute(0, 3, 1, 2) # NHWC -> NCHW
if args.use_lpips_loss:
loss = perception_loss(img, target)
else:
loss = (img - target).pow(2).mean()
print('render loss:', loss.item())
# Backpropagate the gradients.
loss.backward()
# Take a gradient descent step.
points_optim.step()
color_optim.step()
for group in shape_groups:
group.fill_color.data.clamp_(0.0, 1.0)
if t % 10 == 0 or t == args.num_iter - 1:
pydiffvg.save_svg('results/refine_svg/iter_{}.svg'.format(t),
canvas_width, canvas_height, shapes,
shape_groups)
# Render the final result.
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None, # bg
*scene_args)
# Save the intermediate render.
pydiffvg.imwrite(img.cpu(),
'results/refine_svg/final.png'.format(t),
gamma=gamma)
# Convert the intermediate renderings to a video.
from subprocess import call
call([
"ffmpeg", "-framerate", "24", "-i", "results/refine_svg/iter_%d.png",
"-vb", "20M", "results/refine_svg/out.mp4"
])
def save_to_svg(self, name, canvas_height, canvas_width):
pydiffvg.save_svg(f"/content/drive/MyDrive/svgs/{name}_{self.color_name}.svg", canvas_height, canvas_width,
self.color_shapes, self.color_shape_groups, use_gamma=False)
def main(args):
if args.seed:
np.random.seed(args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
pydiffvg.set_print_timing(False)
outdir = os.path.join(args.results_dir, args.prompt, args.subdir)
# Use GPU if available
pydiffvg.set_use_gpu(torch.cuda.is_available())
canvas_width, canvas_height = 224, 224
margin = args.initial_margin
total_paths = args.open_paths + args.closed_paths
step = min(args.step, total_paths)
if step == 0:
step = total_paths
fill_color = None
stroke_color = None
shapes = []
shape_groups = []
losses = []
tt = 0
for num_paths in range(step, total_paths + 1, step):
for i in range(num_paths - step, num_paths):
num_segments = random.randint(1, args.extra_segments + 1)
p0 = (margin + random.random() * (1 - 2 * margin),
margin + random.random() * (1 - 2 * margin))
points = [p0]
is_closed = i >= args.open_paths
if is_closed:
num_segments += 2
for j in range(num_segments):
p1 = (p0[0] + radius * (random.random() - 0.5),
p0[1] + radius * (random.random() - 0.5))
p2 = (p1[0] + radius * (random.random() - 0.5),
p1[1] + radius * (random.random() - 0.5))
p3 = (p2[0] + radius * (random.random() - 0.5),
p2[1] + radius * (random.random() - 0.5))
points.append(p1)
points.append(p2)
if is_closed and j < num_segments - 1 or not is_closed:
points.append(p3)
p0 = p3
points = torch.tensor(points)
points[:, 0] *= canvas_width
points[:, 1] *= canvas_height
stroke_width = torch.tensor(1.0)
color = torch.tensor([
random.random(),
random.random(),
random.random(),
random.random()
])
num_control_points = torch.zeros(num_segments,
dtype=torch.int32) + 2
path = pydiffvg.Path(num_control_points=num_control_points,
points=points,
stroke_width=stroke_width,
is_closed=is_closed)
shapes.append(path)
path_group = pydiffvg.ShapeGroup(
shape_ids=torch.tensor([len(shapes) - 1]),
fill_color=color if is_closed else None,
stroke_color=None if is_closed else color)
shape_groups.append(path_group)
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
render = pydiffvg.RenderFunction.apply
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None,
*scene_args)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
pydiffvg.imwrite(img.cpu(),
os.path.join(outdir, 'init.png'),
gamma=gamma)
points_vars = []
stroke_width_vars = []
color_vars = []
for path in shapes:
path.points.requires_grad = True
points_vars.append(path.points)
if not path.is_closed and args.max_width > 1:
path.stroke_width.requires_grad = True
stroke_width_vars.append(path.stroke_width)
for group in shape_groups:
if group.fill_color is not None:
group.fill_color.requires_grad = True
color_vars.append(group.fill_color)
else:
group.stroke_color.requires_grad = True
color_vars.append(group.stroke_color)
# Embed prompt
text_features = clip_utils.embed_text(args.prompt)
# Optimize
points_optim = torch.optim.Adam(points_vars, lr=args.points_lr)
if len(stroke_width_vars) > 0:
width_optim = torch.optim.Adam(stroke_width_vars, lr=args.width_lr)
color_optim = torch.optim.Adam(color_vars, lr=args.color_lr)
# Adam iterations.
final = False
this_step_iters = max(1, round(args.num_iter * step / total_paths))
if num_paths + step > total_paths:
final = True
this_step_iters += args.extra_iter
for t in range(this_step_iters):
points_optim.zero_grad()
if len(stroke_width_vars) > 0:
width_optim.zero_grad()
color_optim.zero_grad()
# Forward pass: render the image.
scene_args = pydiffvg.RenderFunction.serialize_scene(\
canvas_width, canvas_height, shapes, shape_groups)
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
tt, # seed
None,
*scene_args)
# Save the intermediate render.
with warnings.catch_warnings():
warnings.simplefilter("ignore")
pydiffvg.imwrite(img.cpu(),
os.path.join(outdir,
'iter_{}.png'.format(tt)),
gamma=gamma)
image_features = clip_utils.embed_image(img)
loss = -torch.cosine_similarity(
text_features, image_features, dim=-1).mean()
# Backpropagate the gradients.
loss.backward()
losses.append(loss.item())
# Take a gradient descent step.
points_optim.step()
if len(stroke_width_vars) > 0:
width_optim.step()
color_optim.step()
for path in shapes:
path.points.data[:, 0].clamp_(0.0, canvas_width)
path.points.data[:, 1].clamp_(0.0, canvas_height)
if not path.is_closed:
path.stroke_width.data.clamp_(1.0, args.max_width)
for group in shape_groups:
if group.fill_color is not None:
group.fill_color.data[:3].clamp_(0.0, 1.0)
group.fill_color.data[3].clamp_(args.min_alpha, 1.0)
else:
group.stroke_color.data[:3].clamp_(0.0, 1.0)
group.stroke_color.data[3].clamp_(args.min_alpha, 1.0)
if tt % 10 == 0 or final and t == this_step_iters - 1:
print('%d loss=%.3f' % (tt, 1 + losses[-1]))
pydiffvg.save_svg(
os.path.join(outdir, 'iter_{}.svg'.format(tt)),
canvas_width, canvas_height, shapes, shape_groups)
clip_utils.plot_losses(losses, outdir)
tt += 1
# Render the final result.
img = render(
args.final_px, # width
args.final_px, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None,
*scene_args)
# Save the intermediate render
with warnings.catch_warnings():
warnings.simplefilter("ignore")
pydiffvg.imwrite(img.cpu(),
os.path.join(outdir, 'final.png'),
gamma=gamma)
# Convert the intermediate renderings to a video with a white background.
from subprocess import call
call([
"ffmpeg", "-framerate", "24", "-i",
os.path.join(outdir, "iter_%d.png"), "-vb", "20M", "-filter_complex",
"color=white,format=rgb24[c];[c][0]scale2ref[c][i];[c][i]overlay=format=auto:shortest=1,setsar=1",
"-c:v", "libx264", "-pix_fmt", "yuv420p", "-profile:v", "baseline",
"-movflags", "+faststart",
os.path.join(outdir, "out.mp4")
])
def main(args):
inceptionv3 = models.inception_v3(pretrained=True,
transform_input=False).cuda()
inceptionv3.eval()
perception_loss = ttools.modules.LPIPS().to(pydiffvg.get_device())
canvas_width, canvas_height, shapes, shape_groups = \
pydiffvg.svg_to_scene(args.svg)
scene_args = pydiffvg.RenderFunction.serialize_scene( \
canvas_width, canvas_height, shapes, shape_groups)
render = pydiffvg.RenderFunction.apply
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None, # bg
*scene_args)
# The output image is in linear RGB space. Do Gamma correction before saving the image.
pydiffvg.imwrite(img.cpu(), 'logs/refine_svg/init.png', gamma=gamma)
pydiffvg.imwrite(img.cpu(), 'logs/refine_svg/init_.png')
points_vars = []
for path in shapes:
path.points.requires_grad = True
points_vars.append(path.points)
# color_vars = {}
# for group in shape_groups:
# group.fill_color.requires_grad = True
# color_vars[group.fill_color.data_ptr()] = group.fill_color
# color_vars = list(color_vars.values())
# Optimize
points_optim = torch.optim.Adam(points_vars, lr=1.0)
# color_optim = torch.optim.Adam(color_vars, lr=0.01)
# Adam iterations.
for t in range(args.num_iter):
print('iteration:', t)
points_optim.zero_grad()
# color_optim.zero_grad()
# Forward pass: render the image.
scene_args = pydiffvg.RenderFunction.serialize_scene( \
canvas_width, canvas_height, shapes, shape_groups)
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None, # bg
*scene_args)
# Compose img with white background
img = img[:, :, 3:4] * img[:, :, :3] + torch.ones(
img.shape[0], img.shape[1], 3,
device=pydiffvg.get_device()) * (1 - img[:, :, 3:4])
# Save the intermediate render.
pydiffvg.imwrite(img.cpu(),
'logs/refine_svg/iter_{}.png'.format(t),
gamma=gamma)
img = img[:, :, :3]
# Convert img from HWC to NCHW
img = img.unsqueeze(0)
img = img.permute(0, 3, 1, 2) # NHWC -> NCHW
output = inceptionv3.forward(img.cuda())
get_class(output)
target = torch.autograd.Variable(torch.LongTensor([291]),
requires_grad=False).cuda()
loss = torch.nn.CrossEntropyLoss()(output, target)
print('render loss:', loss.item())
# Backpropagate the gradients.
loss.backward()
# Take a gradient descent step.
points_optim.step()
# color_optim.step()
# for group in shape_groups:
# group.fill_color.data.clamp_(0.0, 1.0)
if t % 10 == 0 or t == args.num_iter - 1:
pydiffvg.save_svg('logs/refine_svg/iter_{}.svg'.format(t),
canvas_width, canvas_height, shapes,
shape_groups)
# Render the final result.
scene_args = pydiffvg.RenderFunction.serialize_scene( \
canvas_width, canvas_height, shapes, shape_groups)
img = render(
canvas_width, # width
canvas_height, # height
2, # num_samples_x
2, # num_samples_y
0, # seed
None, # bg
*scene_args)
# Save the intermediate render.
pydiffvg.imwrite(img.cpu(),
'logs/refine_svg/final.png'.format(t),
gamma=gamma)
def save_svg(self, name):
pydiffvg.save_svg(name, self.canvas_width, self.canvas_height,
self.shapes, self.shape_groups)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--svg",
default=os.path.join("imgs", "seamcarving",
"hokusai.svg"))
parser.add_argument("--optim_steps", default=10, type=int)
parser.add_argument("--lr", default=1e-1, type=int)
args = parser.parse_args()
name = os.path.splitext(os.path.basename(args.svg))[0]
root = os.path.join("results", "seam_carving", name)
svg_root = os.path.join(root, "svg")
os.makedirs(root, exist_ok=True)
os.makedirs(os.path.join(root, "svg"), exist_ok=True)
pydiffvg.set_use_gpu(False)
# pydiffvg.set_device(th.device('cuda'))
# Load SVG
print("loading svg %s" % args.svg)
canvas_width, canvas_height, shapes, shape_groups = \
pydiffvg.svg_to_scene(args.svg)
print("done loading")
max_size = 512
scale_factor = max_size / max(canvas_width, canvas_height)
print("rescaling from %dx%d with scale %f" %
(canvas_width, canvas_height, scale_factor))
canvas_width = int(canvas_width * scale_factor)
canvas_height = int(canvas_height * scale_factor)
print("new shape %dx%d" % (canvas_width, canvas_height))
vector_rescale(shapes, scale_x=scale_factor, scale_y=scale_factor)
# Shrink image by 33 %
# num_seams_to_remove = 2
num_seams_to_remove = canvas_width // 3
new_canvas_width = canvas_width - num_seams_to_remove
scaling = new_canvas_width * 1.0 / canvas_width
# Naive scaling baseline
print("rendering naive rescaling...")
vector_rescale(shapes, scale_x=scaling)
resized = render(new_canvas_width, canvas_height, shapes, shape_groups)
pydiffvg.imwrite(resized.cpu(),
os.path.join(root, 'uniform_scaling.png'),
gamma=2.2)
pydiffvg.save_svg(os.path.join(svg_root, 'uniform_scaling.svg'),
canvas_width,
canvas_height,
shapes,
shape_groups,
use_gamma=False)
vector_rescale(shapes,
scale_x=1.0 / scaling) # bring back original coordinates
print("saved naiving scaling")
# Save initial state
print("rendering initial state...")
im = render(canvas_width, canvas_height, shapes, shape_groups)
pydiffvg.imwrite(im.cpu(), os.path.join(root, 'init.png'), gamma=2.2)
pydiffvg.save_svg(os.path.join(svg_root, 'init.svg'),
canvas_width,
canvas_height,
shapes,
shape_groups,
use_gamma=False)
print("saved initial state")
# Optimize
# color_optim = th.optim.Adam(color_vars, lr=0.01)
retargeted = im[..., :3].cpu().numpy()
previous_width = canvas_width
print("carving seams")
for seam_idx in range(num_seams_to_remove):
print('\nseam', seam_idx + 1, 'of', num_seams_to_remove)
# Remove a seam
retargeted = carve_seam(retargeted)
current_width = canvas_width - seam_idx - 1
scale_factor = current_width * 1.0 / previous_width
previous_width = current_width
padded = np.zeros((canvas_height, canvas_width, 4))
padded[:, :-seam_idx - 1, :3] = retargeted
padded[:, :-seam_idx - 1, -1] = 1.0 # alpha
padded = th.from_numpy(padded).to(im.device)
# Remap points to the smaller canvas and
# collect variables to optimize
points_vars = []
# width_vars = []
mini, maxi = canvas_width, 0
for path in shapes:
path.points.requires_grad = False
x = path.points[..., 0]
y = path.points[..., 1]
# rescale
x = x * scale_factor
# clip to canvas
path.points[..., 0] = th.clamp(x, 0, current_width)
path.points[..., 1] = th.clamp(y, 0, canvas_height)
path.points.requires_grad = True
points_vars.append(path.points)
path.stroke_width.requires_grad = True
# width_vars.append(path.stroke_width)
mini = min(mini, path.points.min().item())
maxi = max(maxi, path.points.max().item())
print("points", mini, maxi, "scale", scale_factor)
# recreate an optimizer so we don't carry over the previous update
# (momentum)?
geom_optim = th.optim.Adam(points_vars, lr=args.lr)
for step in range(args.optim_steps):
geom_optim.zero_grad()
img = render(canvas_width,
canvas_height,
shapes,
shape_groups,
samples=2)
pydiffvg.imwrite(img.cpu(),
os.path.join(
root,
"seam_%03d_iter_%02d.png" % (seam_idx, step)),
gamma=2.2)
# NO alpha
loss = (img - padded)[..., :3].pow(2).mean()
# loss = (img - padded).pow(2).mean()
print('render loss:', loss.item())
# Backpropagate the gradients.
loss.backward()
# Take a gradient descent step.
geom_optim.step()
pydiffvg.save_svg(os.path.join(svg_root, "seam%03d.svg" % seam_idx),
canvas_width - seam_idx,
canvas_height,
shapes,
shape_groups,
use_gamma=False)
for path in shapes:
mini = min(mini, path.points.min().item())
maxi = max(maxi, path.points.max().item())
print("points", mini, maxi)
img = render(canvas_width, canvas_height, shapes, shape_groups)
img = img[:, :-num_seams_to_remove]
pydiffvg.imwrite(img.cpu(), os.path.join(root, 'final.png'), gamma=2.2)
pydiffvg.imwrite(retargeted, os.path.join(root, 'ref.png'), gamma=2.2)
pydiffvg.save_svg(os.path.join(svg_root, 'final.svg'),
canvas_width - num_seams_to_remove + 1,
canvas_height,
shapes,
shape_groups,
use_gamma=False)
# Convert the intermediate renderings to a video.
from subprocess import call
call([
"ffmpeg", "-framerate", "24", "-i",
os.path.join(root, "seam_%03d_iter_00.png"), "-vb", "20M",
os.path.join(root, "out.mp4")
])
def generate_samples(args):
chkpt = VAE_OUTPUT
if args.conditional:
chkpt += "_conditional"
if args.fc:
chkpt += "_fc"
meta = ttools.Checkpointer.load_meta(chkpt, prefix="g_")
if meta is None:
LOG.info("No metadata in checkpoint (or no checkpoint), aborting.")
return
model = VectorMNISTVAE(**meta)
checkpointer = ttools.Checkpointer(chkpt, model, prefix="g_")
checkpointer.load_latest()
model.eval()
# Sample some latent vectors
n = 8
bs = n*n
z = th.randn(bs, model.zdim)
imsize = 28
dataset = Dataset(args.data_dir, imsize)
dataloader = DataLoader(dataset, batch_size=bs,
num_workers=1, shuffle=True)
for batch in dataloader:
ref, label = batch
break
autoencode = True
if autoencode:
LOG.info("Sampling with auto-encoder code")
if not args.conditional:
label = None
mu, logvar = model.encode(ref, label)
z = model.reparameterize(mu, logvar)
else:
label = None
if args.conditional:
label = th.clamp(th.rand(bs)*10, 0, 9).long()
if args.digit is not None:
label[:] = args.digit
with th.no_grad():
images, aux = model.decode(z, label=label)
scenes = aux["scenes"]
images += 1.0
images /= 2.0
h = w = model.imsize
images = images.view(n, n, h, w).permute(0, 2, 1, 3)
images = images.contiguous().view(n*h, n*w)
images = th.clamp(images, 0, 1).cpu().numpy()
path = os.path.join(chkpt, "samples.png")
pydiffvg.imwrite(images, path, gamma=2.2)
if autoencode:
ref += 1.0
ref /= 2.0
ref = ref.view(n, n, h, w).permute(0, 2, 1, 3)
ref = ref.contiguous().view(n*h, n*w)
ref = th.clamp(ref, 0, 1).cpu().numpy()
path = os.path.join(chkpt, "ref.png")
pydiffvg.imwrite(ref, path, gamma=2.2)
# merge scenes
all_shapes = []
all_shape_groups = []
cur_id = 0
for idx, s in enumerate(scenes):
shapes, shape_groups, _ = s
# width, height = sizes
# Shift digit on canvas
center_x = idx % n
center_y = idx // n
for shape in shapes:
shape.points[:, 0] += center_x * model.imsize
shape.points[:, 1] += center_y * model.imsize
all_shapes.append(shape)
for grp in shape_groups:
grp.shape_ids[:] = cur_id
cur_id += 1
all_shape_groups.append(grp)
LOG.info("Generated %d shapes", len(all_shapes))
fname = os.path.join(chkpt, "digits.svg")
pydiffvg.save_svg(fname, n*model.imsize, n*model.imsize, all_shapes,
all_shape_groups, use_gamma=False)
LOG.info("Results saved to %s", chkpt)
def save(self,
all_points,
save_dir,
name,
verbose=False,
white_background=True):
# note that this if for a single shape and bs dimension should have multiple curves
# print('1:', process.memory_info().rss*1e-6)
render_size = self.imsize
bs = all_points.shape[0]
if verbose:
render_size = render_size * 2
all_points = all_points * render_size
num_ctrl_pts = torch.zeros(self.curves, dtype=torch.int32) + 2
shapes = []
shape_groups = []
for k in range(bs):
# Get point parameters from network
color = make_tensor(color[k])
points = all_points[k].cpu().contiguous() #[self.sort_idx[k]]
if verbose:
np.random.seed(0)
colors = np.random.rand(self.curves, 4)
high = np.array((0.565, 0.392, 0.173, 1))
low = np.array((0.094, 0.310, 0.635, 1))
diff = (high - low) / (self.curves)
colors[:, 3] = 1
for i in range(self.curves):
scale = diff * i
color = low + scale
color[3] = 1
color = torch.tensor(color)
num_ctrl_pts = torch.zeros(1, dtype=torch.int32) + 2
if i * 3 + 4 > self.curves * 3:
curve_points = torch.stack([
points[i * 3], points[i * 3 + 1],
points[i * 3 + 2], points[0]
])
else:
curve_points = points[i * 3:i * 3 + 4]
path = pydiffvg.Path(num_control_points=num_ctrl_pts,
points=curve_points,
is_closed=False,
stroke_width=torch.tensor(4))
path_group = pydiffvg.ShapeGroup(shape_ids=torch.tensor(
[i]),
fill_color=None,
stroke_color=color)
shapes.append(path)
shape_groups.append(path_group)
for i in range(self.curves * 3):
scale = diff * (i // 3)
color = low + scale
color[3] = 1
color = torch.tensor(color)
if i % 3 == 0:
# color = torch.tensor(colors[i//3]) #green
shape = pydiffvg.Rect(p_min=points[i] - 8,
p_max=points[i] + 8)
group = pydiffvg.ShapeGroup(shape_ids=torch.tensor(
[self.curves + i]),
fill_color=color)
else:
# color = torch.tensor(colors[i//3]) #purple
shape = pydiffvg.Circle(radius=torch.tensor(8.0),
center=points[i])
group = pydiffvg.ShapeGroup(shape_ids=torch.tensor(
[self.curves + i]),
fill_color=color)
shapes.append(shape)
shape_groups.append(group)
else:
path = pydiffvg.Path(num_control_points=num_ctrl_pts,
points=points,
is_closed=True)
shapes.append(path)
path_group = pydiffvg.ShapeGroup(shape_ids=torch.tensor(
[len(shapes) - 1]),
fill_color=color,
stroke_color=color)
shape_groups.append(path_group)
pydiffvg.save_svg(f"{save_dir}{name}/{name}.svg", self.imsize,
self.imsize, shapes, shape_groups)
def scene_to_svg(scene, path):
canvas_width, canvas_height, shapes, shape_groups = scene
pydiffvg.save_svg(path,
canvas_width, canvas_height, shapes, shape_groups)
