94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
"""
|
|
"""
|
|
import os
|
|
import pydiffvg
|
|
import torch as th
|
|
import scipy.ndimage.filters as F
|
|
|
|
|
|
def render(canvas_width, canvas_height, shapes, shape_groups):
|
|
_render = pydiffvg.RenderFunction.apply
|
|
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,
|
|
*scene_args)
|
|
return img
|
|
|
|
|
|
def main():
|
|
pydiffvg.set_device(th.device('cuda:1'))
|
|
|
|
# Load SVG
|
|
svg = os.path.join("imgs", "peppers.svg")
|
|
canvas_width, canvas_height, shapes, shape_groups = \
|
|
pydiffvg.svg_to_scene(svg)
|
|
|
|
# Save initial state
|
|
ref = render(canvas_width, canvas_height, shapes, shape_groups)
|
|
pydiffvg.imwrite(ref.cpu(), 'results/gaussian_blur/init.png', gamma=2.2)
|
|
|
|
target = F.gaussian_filter(ref.cpu().numpy(), [10, 10, 0])
|
|
target = th.from_numpy(target).to(ref.device)
|
|
pydiffvg.imwrite(target.cpu(), 'results/gaussian_blur/target.png', gamma=2.2)
|
|
|
|
# Collect variables to optimize
|
|
points_vars = []
|
|
width_vars = []
|
|
for path in shapes:
|
|
path.points.requires_grad = True
|
|
points_vars.append(path.points)
|
|
path.stroke_width.requires_grad = True
|
|
width_vars.append(path.stroke_width)
|
|
color_vars = []
|
|
for group in shape_groups:
|
|
# do not optimize alpha
|
|
group.fill_color[..., :3].requires_grad = True
|
|
color_vars.append(group.fill_color)
|
|
|
|
# Optimize
|
|
points_optim = th.optim.Adam(points_vars, lr=1.0)
|
|
width_optim = th.optim.Adam(width_vars, lr=1.0)
|
|
color_optim = th.optim.Adam(color_vars, lr=0.01)
|
|
|
|
for t in range(20):
|
|
print('\niteration:', t)
|
|
points_optim.zero_grad()
|
|
width_optim.zero_grad()
|
|
color_optim.zero_grad()
|
|
# Forward pass: render the image.
|
|
img = render(canvas_width, canvas_height, shapes, shape_groups)
|
|
# Save the intermediate render.
|
|
pydiffvg.imwrite(img.cpu(), 'results/gaussian_blur/iter_{}.png'.format(t), gamma=2.2)
|
|
loss = (img - target)[..., :3].pow(2).mean()
|
|
|
|
print('alpha:', img[..., 3].mean().item())
|
|
print('render loss:', loss.item())
|
|
|
|
# Backpropagate the gradients.
|
|
loss.backward()
|
|
|
|
# Take a gradient descent step.
|
|
points_optim.step()
|
|
width_optim.step()
|
|
color_optim.step()
|
|
for group in shape_groups:
|
|
group.fill_color.data.clamp_(0.0, 1.0)
|
|
|
|
# Final render
|
|
img = render(canvas_width, canvas_height, shapes, shape_groups)
|
|
pydiffvg.imwrite(img.cpu(), 'results/gaussian_blur/final.png', gamma=2.2)
|
|
|
|
# Convert the intermediate renderings to a video.
|
|
from subprocess import call
|
|
call(["ffmpeg", "-framerate", "24", "-i",
|
|
"results/gaussian_blur/iter_%d.png", "-vb", "20M",
|
|
"results/gaussian_blur/out.mp4"])
|
|
|
|
if __name__ == "__main__":
|
|
main()
|