refinement.py

import pydiffvg
import argparse
import torch
import skimage.io
import os
import re
from shutil import copyfile
from PIL import Image
import numpy as np
gamma = 1.0

def cal_alignment_loss(args, save_path):
    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.
    points_vars = []
    for path in shapes:
        #print(path)
        #input()
        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)
    color_optim = torch.optim.Adam(color_vars, lr=0)

    # Adam iterations.
    for t in range(args.num_iter):
        
        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])
        img = img[:, :, :3]
        # Convert img from HWC to NCHW
        img = img.unsqueeze(0)
        img = img.permute(0, 3, 1, 2) # NHWC -> NCHW
        loss = (img - target).pow(2).mean()
        #if t%10 == 0:
        #    print('iteration:', t)
        #    print('render loss:', args.no_sample, 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 == args.num_iter - 1:

            pydiffvg.save_svg_paths_only(save_path, canvas_width, canvas_height, shapes, shape_groups)
        
    return loss

def get_svg_glyph_bbox(svg_path):
    fin = open(svg_path,'r')
    path_ = fin.read().split('d="')[1]
    path = path_.split('" fill=')[0]
    path_splited = re.split(r"([mlc])", path)
    commands = []
    cur_x = 0.0
    cur_y = 0.0
    x_min = 1000
    x_max = -1000
    y_min = 1000
    y_max = -1000
    first_move = True
    for idx in range(0,len(path_splited)):
        if len(path_splited[idx]) == 0: continue
        # x1,y1,x2,y2,x3,y3,x4,y4 are the absolute coords
        if path_splited[idx] == 'm':
            coords_str = path_splited[idx+1]
            print(first_move)
            if first_move:
                x4 = float(coords_str.split(' ')[1])
                y4 = float(coords_str.split(' ')[2])
                first_move = False
            else:
                x4 = cur_x + float(coords_str.split(' ')[1])
                y4 = cur_y + float(coords_str.split(' ')[2])
                               
            cur_x = x4
            cur_y = y4
            x_min = min(cur_x, x_min)
            x_max = max(cur_x, x_max)
            y_min = min(cur_y, y_min)
            y_max = max(cur_y, y_max)
            print(cur_x,cur_y)
        if path_splited[idx] == 'l':
            coords_str = path_splited[idx+1]
            x4 = cur_x + float(coords_str.split(' ')[1])
            y4 = cur_y + float(coords_str.split(' ')[2])
            cur_x = x4
            cur_y = y4
            x_min = min(cur_x, x_min)
            x_max = max(cur_x, x_max)
            y_min = min(cur_y, y_min)
            y_max = max(cur_y, y_max)
            print(cur_x,cur_y)
        if path_splited[idx] == 'c':
            coords_str = path_splited[idx+1]
            x1 = cur_x
            y1 = cur_y
            x2 = cur_x + float(coords_str.split(' ')[1])
            y2 = cur_y + float(coords_str.split(' ')[2])
            x3 = cur_x + float(coords_str.split(' ')[3])
            y3 = cur_y + float(coords_str.split(' ')[4])
            x4 = cur_x + float(coords_str.split(' ')[5])
            y4 = cur_y + float(coords_str.split(' ')[6])
            x_min = min(x2, x3, x4, x_min)
            x_max = max(x2, x3, x4, x_max)
            y_min = min(y2, y3, y4, y_min)
            y_max = max(y2, y3, y4, y_max)
            cur_x = x4
            cur_y = y4
            print(cur_x,cur_y)
    return [x_min,x_max], [y_min,y_max]

def get_img_bbox(img_path):
    print(img_path)
    img = Image.open(img_path)
    img = 255 - np.array(img)
    img0 = np.sum(img, axis = 0)
    img1 = np.sum(img, axis = 1)
    y_range = np.where(img1>127.5)[0]
    x_range = np.where(img0>127.5)[0]
    return [x_range[0],x_range[-1]], [y_range[0],y_range[-1]]

def svg_bbox_align(svg_path, trgimg_path):

    svg_xr, svg_yr = get_svg_glyph_bbox(svg_path)
    img_xr, img_yr = get_img_bbox(trgimg_path)
    svg_w = svg_xr[1] - svg_xr[0]
    svg_h = svg_yr[1] - svg_yr[0]
    svg_xc = (svg_xr[1] + svg_xr[0]) / 2.0
    svg_yc = (svg_yr[1] + svg_yr[0]) / 2.0
    img_w = img_xr[1] - img_xr[0] + 1
    img_h = img_yr[1] - img_yr[0] + 1
    img_xc = (img_xr[1] + img_xr[0]) / 2.0
    img_yc = (img_yr[1] + img_yr[0]) / 2.0    

    def affine_coord(coord, x_or_y, cur_cmd, first_move):
        
        if x_or_y % 2 == 0: # for x
            if cur_cmd == 'm' and first_move:
                new_coord = (coord - svg_xc) * (img_w / svg_w) + img_xc
                res = str(new_coord)
            else:
                res = str((img_w / svg_w) * (coord))
        else: # for y 
            if cur_cmd == 'm' and first_move:
                new_coord = (coord - svg_yc) * (img_h / svg_h) + img_yc
                res = str(new_coord)
            else:
                res = str((img_h / svg_h) * (coord))
    
        return res

    svg_raw = open(svg_path,'r').read()
    fout = open(svg_path.split('.svg')[0] + '_256.svg','w')
    fout.write('<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="256px" height="256px" style="-ms-transform: rotate(360deg); -webkit-transform: rotate(360deg); transform: rotate(360deg);" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256">')
    
    coord = '<path' + svg_raw.split('<path')[1]
    tokens = coord.split(' ')
    newcoord = ''
    first_move = True
    x_or_y = 0
    for k in tokens:
        if k[0] != '<' and k[0] != 'd' and k[0] != 'm' and k[0] != 'c' and k[0] != 'l' and k[0] != 'f':
            if k[-1] != '"':
                newcoord += affine_coord(float(k), x_or_y, cur_cmd, first_move)
                if cur_cmd == 'm': first_move = False
                x_or_y += 1
                newcoord += ' '
            else:
                newcoord += affine_coord(float(k[0:len(k)-1]), x_or_y, cur_cmd, first_move)
                x_or_y += 1                             
                newcoord += '" '
        else:
            cur_cmd = k
            newcoord += k
            newcoord += ' '
    fout.write(newcoord)
    fout.close()



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--svg", help="source SVG path", type=str, default='none')
    parser.add_argument("--target", help="target image path", type=str, default='none')
    parser.add_argument("--use_lpips_loss", dest='use_lpips_loss', action='store_true')
    parser.add_argument("--num_iter", type=int, default=40)
    parser.add_argument("--no_sample", type=int, default=0)
    parser.add_argument("--fontid", type=str, default='0')
    parser.add_argument("--experiment_name", type=str, default='v1.0_gumbletrain')
    parser.add_argument("--candidate_nums", type=str, default='20')
    args = parser.parse_args()

    svg_path = os.path.join('experiments', args.experiment_name + '_main_model/results/', '%04d'%int(args.fontid), 'svgs')
    imghr_path = os.path.join('experiments', args.experiment_name + '_main_model/results/', '%04d'%int(args.fontid), 'imgs_256')
    svg_outpath = os.path.join('experiments', args.experiment_name + '_main_model/results/', '%04d'%int(args.fontid), 'svgs-refined')
    if not os.path.exists(svg_outpath):
        os.mkdir(svg_outpath)
    

    for i in range(0,52):#62

        # find the best candidate
        minLoss = 10000
        noMin = 0
        tempLoss = 0
        # pick the best candidate
        for j in range(0, int(args.candidate_nums)):
            print(f'processing_char_{i:02d}_candidate_{j:02d}')
            args.no_sample = j
            args.svg = os.path.join(svg_path, 'syn_%02d_%02d.svg'%(i,j))
            args.target = os.path.join(imghr_path, '%02d_256.png'%i)
            #svg_aligned = align(args.svg, args.target)
            svg_bbox_align(args.svg, args.target)
            args.svg = os.path.join(svg_path, 'syn_%02d_%02d_256.svg'%(i,j))
            #svg_init_aligned = os.path.join(svg_path, 'syn_%02d_'%i, '%02d'%j, '.svg')
            tempLoss = cal_alignment_loss(args, save_path = args.svg.split('.svg')[0] + '_r.svg')
            #print(f'finished_char_{i:02d}_candidate_{j:02d}')
            if tempLoss < minLoss:
                noMin = j
                minLoss = tempLoss
        # do longer optimization
        args.num_iter = 300
        args.svg = os.path.join(svg_path, 'syn_%02d_%02d_256_r.svg'%(i,noMin))
        tempLoss = cal_alignment_loss(args, save_path = os.path.join(svg_outpath, 'syn_%02d.svg'%(i)))

    svg_merge_outpath = os.path.join(svg_outpath, f"syn_svg_merge.html")
    fout = open(svg_merge_outpath, 'w')
    for i in range(0,52):
        svg = open(os.path.join(svg_outpath, 'syn_%02d.svg'%(i)),'r')
        f.write(svg)
        if i > 0 and i % 13 == 12:
            fout.write('<br>')
    fout.close()