initial commit

apps/svg_brush.py

@@ -0,0 +1,167 @@
+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()
+