initial commit

apps/texture_synthesis.py

@@ -0,0 +1,180 @@
+import os, sys
+import pydiffvg
+import argparse
+import torch
+# import torch as th
+import scipy.ndimage.filters as filters
+# import numba
+import numpy as np
+from skimage import io
+sys.path.append('./textureSyn')
+from patchBasedTextureSynthesis import *
+from make_gif import make_gif
+import random
+import ttools.modules
+
+from svgpathtools import svg2paths2, Path, is_path_segment
+"""
+python texture_synthesis.py textureSyn/traced_1.png  --svg-path textureSyn/traced_1.svg --case 1
+"""
+
+def texture_syn(img_path):
+    ## get the width and height first
+    # input_img = io.imread(img_path)  # returns an MxNx3 array
+    # output_size = [input_img.shape[1], input_img.shape[0]]
+    # output_path = "textureSyn/1/"
+    output_path = "results/texture_synthesis/%d"%(args.case)
+    patch_size = 40  # size of the patch (without the overlap)
+    overlap_size = 10  # the width of the overlap region
+    output_size = [300, 300]
+    pbts = patchBasedTextureSynthesis(img_path, output_path, output_size, patch_size, overlap_size, in_windowStep=5,
+                                      in_mirror_hor=True, in_mirror_vert=True, in_shapshots=False)
+    target_img = pbts.resolveAll()
+    return np.array(target_img)
+
+
+def render(canvas_width, canvas_height, shapes, shape_groups, samples=2):
+    _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
+                 samples,   # num_samples_x
+                 samples,   # num_samples_y
+                 0,   # seed
+                 None,
+                 *scene_args)
+    return img
+
+def big_bounding_box(paths_n_stuff):
+    """Finds a BB containing a collection of paths, Bezier path segments, and
+    points (given as complex numbers)."""
+    bbs = []
+    for thing in paths_n_stuff:
+        if is_path_segment(thing) or isinstance(thing, Path):
+            bbs.append(thing.bbox())
+        elif isinstance(thing, complex):
+            bbs.append((thing.real, thing.real, thing.imag, thing.imag))
+        else:
+            try:
+                complexthing = complex(thing)
+                bbs.append((complexthing.real, complexthing.real,
+                            complexthing.imag, complexthing.imag))
+            except ValueError:
+                raise TypeError(
+                    "paths_n_stuff can only contains Path, CubicBezier, "
+                    "QuadraticBezier, Line, and complex objects.")
+    xmins, xmaxs, ymins, ymaxs = list(zip(*bbs))
+    xmin = min(xmins)
+    xmax = max(xmaxs)
+    ymin = min(ymins)
+    ymax = max(ymaxs)
+    return xmin, xmax, ymin, ymax
+
+
+def main(args):
+    ## set device -> use cpu now since I haven't solved the nvcc issue
+    pydiffvg.set_use_gpu(False)
+    # pydiffvg.set_device(torch.device('cuda:1'))
+    ## use L2 for now
+    # perception_loss = ttools.modules.LPIPS().to(pydiffvg.get_device())
+
+    ## generate a texture synthesized
+    target_img = texture_syn(args.target)
+    tar_h, tar_w = target_img.shape[1], target_img.shape[0]
+    canvas_width, canvas_height, shapes, shape_groups = \
+        pydiffvg.svg_to_scene(args.svg_path)
+
+
+    ## svgpathtools for checking the bounding box
+    # paths, _, _ = svg2paths2(args.svg_path)
+    # print(len(paths))
+    # xmin, xmax, ymin, ymax = big_bounding_box(paths)
+    # print(xmin, xmax, ymin, ymax)
+    # input("check")
+
+
+    print('tar h : %d tar w : %d'%(tar_h, tar_w))
+    print('canvas h : %d canvas w : %d' % (canvas_height, canvas_width))
+    scale_ratio = tar_h / canvas_height
+    print("scale ratio : ", scale_ratio)
+    # input("check")
+    for path in shapes:
+        path.points[..., 0] = path.points[..., 0] * scale_ratio
+        path.points[..., 1] = path.points[..., 1] * scale_ratio
+
+    init_img = render(tar_w, tar_h, shapes, shape_groups)
+    pydiffvg.imwrite(init_img.cpu(), 'results/texture_synthesis/%d/init.png'%(args.case), gamma=2.2)
+    # input("check")
+    random.seed(1234)
+    torch.manual_seed(1234)
+
+    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.append(group.fill_color)
+    # Optimize
+    points_optim = torch.optim.Adam(points_vars, lr=1.0)
+    color_optim = torch.optim.Adam(color_vars, lr=0.01)
+
+    target = torch.from_numpy(target_img).to(torch.float32) / 255.0
+    target = target.pow(2.2)
+    target = target.to(pydiffvg.get_device())
+    target = target.unsqueeze(0)
+    target = target.permute(0, 3, 1, 2) # NHWC -> NCHW
+    canvas_width, canvas_height = target.shape[3], target.shape[2]
+    # print('canvas h : %d canvas w : %d' % (canvas_height, canvas_width))
+    # input("check")
+
+    for t in range(args.max_iter):
+        print('iteration:', t)
+        points_optim.zero_grad()
+        color_optim.zero_grad()
+        cur_img = render(canvas_width, canvas_height, shapes, shape_groups)
+        pydiffvg.imwrite(cur_img.cpu(), 'results/texture_synthesis/%d/iter_%d.png'%(args.case, t), gamma=2.2)
+        cur_img = cur_img[:, :, :3]
+        cur_img = cur_img.unsqueeze(0)
+        cur_img = cur_img.permute(0, 3, 1, 2) # NHWC -> NCHW
+
+        ## perceptual loss
+        # loss = perception_loss(cur_img, target)
+        ## l2 loss
+        loss = (cur_img - target).pow(2).mean()
+        print('render loss:', loss.item())
+        loss.backward()
+
+        points_optim.step()
+        color_optim.step()
+
+        for group in shape_groups:
+            group.fill_color.data.clamp_(0.0, 1.0)
+        ## write svg
+        if t % 10 == 0 or t == args.max_iter - 1:
+            pydiffvg.save_svg('results/texture_synthesis/%d/iter_%d.svg'%(args.case, t),
+                              canvas_width, canvas_height, shapes, shape_groups)
+
+    ## render final result
+    final_img = render(tar_h, tar_w, shapes, shape_groups)
+    pydiffvg.imwrite(final_img.cpu(), 'results/texture_synthesis/%d/final.png'%(args.case), gamma=2.2)
+
+
+    from subprocess import call
+    call(["ffmpeg", "-framerate", "24", "-i",
+        "results/texture_synthesis/%d/iter_%d.png"%(args.case), "-vb", "20M",
+        "results/texture_synthesis/%d/out.mp4"%(args.case)])
+    ## make gif
+    make_gif("results/texture_synthesis/%d"%(args.case), "results/texture_synthesis/%d/out.gif"%(args.case), frame_every_X_steps=1, repeat_ending=3, total_iter=args.max_iter)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    ## target image path
+    parser.add_argument("target", help="target image path")
+    parser.add_argument("--svg-path", type=str, help="the corresponding svg file path")
+    parser.add_argument("--max-iter", type=int, default=500, help="the max optimization iterations")
+    parser.add_argument("--case", type=int, default=1, help="just the case id for a separate result folder")
+    args = parser.parse_args()
+    main(args)