168 lines
4.9 KiB
Python
168 lines
4.9 KiB
Python
import sys
|
|
sys.path.append("../svg")
|
|
from geometry import GeometryLoss
|
|
import numpy as np
|
|
import pygame as pg
|
|
import torch
|
|
import pydiffvg
|
|
import tkinter as tk
|
|
from tkinter import filedialog
|
|
|
|
def box_kernel(val):
|
|
return np.heaviside(-val+1,0)
|
|
|
|
def cone_kernel(val):
|
|
return np.maximum(0,1-val)
|
|
|
|
def nptosurf(arr):
|
|
if arr.shape[2]==1:
|
|
#greyscale
|
|
shape=arr.shape
|
|
shape=(shape[0],shape[1],3)
|
|
arr=np.broadcast_to(arr,shape)
|
|
return pg.surfarray.make_surface(arr*255)
|
|
|
|
def brush_tensor(screen_size,coords,radius,kernel):
|
|
coordarr=np.stack(np.meshgrid(np.linspace(0,screen_size[0]-1,screen_size[0]),np.linspace(0,screen_size[1]-1,screen_size[1]),indexing='ij'),axis=2)
|
|
ctrarr = np.reshape(np.array(coords), [1, 1, 2])
|
|
distarr=np.sqrt(np.sum(np.power(coordarr-ctrarr,2),axis=2))
|
|
valarr=kernel(distarr/radius)
|
|
return torch.tensor(valarr,requires_grad=False,dtype=torch.float32)
|
|
|
|
def checkerboard(shape, square_size=2):
|
|
xv,yv=np.meshgrid(np.floor(np.linspace(0,shape[1]-1,shape[1])/square_size),np.floor(np.linspace(0,shape[0]-1,shape[0])/square_size))
|
|
bin=np.expand_dims(((xv+yv)%2),axis=2)
|
|
res=bin*np.array([[[1., 1., 1.,]]])+(1-bin)*np.array([[[.75, .75, .75,]]])
|
|
return torch.tensor(res,requires_grad=False,dtype=torch.float32)
|
|
|
|
def render(optim, viewport):
|
|
scene_args = pydiffvg.RenderFunction.serialize_scene(*optim.build_scene())
|
|
render = pydiffvg.RenderFunction.apply
|
|
img = render(viewport[0], # width
|
|
viewport[1], # height
|
|
2, # num_samples_x
|
|
2, # num_samples_y
|
|
0, # seed
|
|
None,
|
|
*scene_args)
|
|
return img
|
|
|
|
def optimize(optim, viewport, brush_kernel, increase=True, strength=0.1):
|
|
optim.zero_grad()
|
|
|
|
geomLoss=torch.tensor(0.)
|
|
|
|
for shape, gloss in zip(optim.scene[2],geometryLosses):
|
|
geomLoss+=gloss.compute(shape)
|
|
|
|
img=render(optim,viewport)
|
|
|
|
imalpha=img[:,:,3]
|
|
|
|
multiplied=imalpha*brush_kernel
|
|
|
|
loss=((1-multiplied).mean() if increase else multiplied.mean())*strength
|
|
|
|
loss+=geomLoss
|
|
|
|
loss.backward()
|
|
|
|
optim.step()
|
|
|
|
return render(optim,viewport)
|
|
|
|
def get_infile():
|
|
pydiffvg.set_use_gpu(False)
|
|
root = tk.Tk()
|
|
#root.withdraw()
|
|
|
|
file_path = filedialog.askopenfilename(initialdir = ".",title = "Select graphic to optimize",filetypes = (("SVG files","*.svg"),("all files","*.*")))
|
|
|
|
root.destroy()
|
|
|
|
return file_path
|
|
|
|
def compositebg(img):
|
|
bg=checkerboard(img.shape,2)
|
|
color=img[:,:,0:3]
|
|
alpha=img[:,:,3]
|
|
composite=alpha.unsqueeze(2)*color+(1-alpha).unsqueeze(2)*bg
|
|
|
|
return composite
|
|
|
|
def main():
|
|
infile=get_infile()
|
|
|
|
settings=pydiffvg.SvgOptimizationSettings()
|
|
settings.global_override(["optimize_color"],False)
|
|
settings.global_override(["transforms","optimize_transforms"], False)
|
|
settings.global_override(["optimizer"], "SGD")
|
|
settings.global_override(["paths","shape_lr"], 1e-1)
|
|
|
|
optim=pydiffvg.OptimizableSvg(infile,settings)
|
|
|
|
global geometryLosses
|
|
geometryLosses = []
|
|
|
|
for shape in optim.build_scene()[2]:
|
|
geometryLosses.append(GeometryLoss(shape))
|
|
|
|
scaling=1
|
|
brush_radius=100
|
|
graphic_size=optim.canvas
|
|
screen_size=(graphic_size[1]*scaling, graphic_size[0]*scaling)
|
|
|
|
pg.init()
|
|
|
|
screen=pg.display.set_mode(screen_size)
|
|
screen.fill((255,255,255))
|
|
|
|
img=render(optim,graphic_size)
|
|
print(img.max())
|
|
|
|
npsurf = pg.transform.scale(nptosurf(compositebg(img).detach().permute(1,0,2).numpy()), screen_size)
|
|
|
|
screen.blit(npsurf,(0,0))
|
|
|
|
pg.display.update()
|
|
clock=pg.time.Clock()
|
|
|
|
z=0
|
|
btn=0
|
|
|
|
while True:
|
|
clock.tick(60)
|
|
for event in pg.event.get():
|
|
if event.type==pg.QUIT:
|
|
pg.quit()
|
|
sys.exit()
|
|
|
|
y, x = pg.mouse.get_pos()
|
|
if event.type == pg.MOUSEBUTTONDOWN:
|
|
if event.button in [1,3]:
|
|
z=1
|
|
btn=event.button
|
|
elif event.button == 4:
|
|
brush_radius*=1.1
|
|
elif event.button == 5:
|
|
brush_radius/=1.1
|
|
brush_radius=max(brush_radius,5)
|
|
elif event.type == pg.MOUSEBUTTONUP:
|
|
if event.button in [1,3]:
|
|
z=0
|
|
|
|
if z==1:
|
|
brush=brush_tensor((graphic_size[0],graphic_size[1]), (x/scaling, y/scaling), brush_radius, box_kernel)
|
|
img=optimize(optim,graphic_size,brush,btn==1)
|
|
npsurf = pg.transform.scale(nptosurf(compositebg(img).detach().permute(1,0,2).numpy()), screen_size)
|
|
|
|
|
|
screen.blit(npsurf,(0,0))
|
|
pg.draw.circle(screen, (255,255,255), (y,x), int(brush_radius*scaling), 1)
|
|
pg.display.update()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|