many changes

gradientmesh/test.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+from time import sleep
+
+import pydiffvg
+import torch
+import random
+from random import uniform
+
+from gmtypes import GradientMesh, Quad, Patch, Point, join_quads
+
+
+def quads():
+    return [
+        Quad.random(),
+        Quad.random(),
+        Quad.random(),
+        Quad.random(),
+    ]
+
+
+def rand_quad_test(filename='random_quad.png', width=256, height=256,
+                   degree=4, num_control_points=2):
+    pydiffvg.set_use_gpu(torch.cuda.is_available())
+    render = pydiffvg.RenderFunction.apply
+
+    patch = Patch.random()
+
+    shape_groups = [patch.as_shape_group()]
+    shapes = [patch.as_path(width, height)]
+
+    scene_args = pydiffvg.RenderFunction.serialize_scene(width, height,
+                                                         shapes, shape_groups)
+
+    img = render(width, height, 2, 2, 0, None, *scene_args)
+    pydiffvg.imwrite(img.cpu(), f"test_data/{filename}", gamma=2.2)
+
+    return img
+
+
+def mult_quad_test(filename='multiple_quads.png', width=1024,
+                   height=1024, num_control_points=None, mask=None, seed=None):
+    random.seed(seed)
+    mask = mask or [1, 1, 1, 1]
+    pydiffvg.set_use_gpu(torch.cuda.is_available())
+    render = pydiffvg.RenderFunction.apply
+
+    a, b, c, d = quads()
+    join_quads(a, b, c, d)
+
+    to_render = [a, b, c, d]
+    to_render = [x for x in to_render if mask[to_render.index(x)]]
+
+    shape_groups = [patch.as_shape_group(color=(
+        uniform(0, 1),
+        uniform(0, 1),
+        uniform(0, 1),
+        0.8
+    )) for patch in to_render]
+
+    for i in range(len(to_render)):
+        shape_groups[i].shape_ids = torch.tensor([i])
+    shapes = [patch.as_path(width, height) for patch in to_render]
+
+    scene_args = pydiffvg.RenderFunction.serialize_scene(width, height,
+                                                         shapes, shape_groups)
+
+    img = render(width, height, 2, 2, 0, None, *scene_args)
+    pydiffvg.imwrite(img.cpu(), f"test_data/{filename}", gamma=2.2)
+    return img.clone()
+
+
+def om():
+    filename = 'optimize_test.png'
+    pydiffvg.set_use_gpu(torch.cuda.is_available())
+    render = pydiffvg.RenderFunction.apply
+
+    target = mult_quad_test(width=256, height=256)
+
+    squad = quads()
+
+    join_quads(*squad)
+
+    gm = GradientMesh(*squad)
+
+    points_n = []
+    for s in squad:
+        out = []
+        for pt in s.points:
+            out.append([pt.x, pt.y])
+            for cpt in pt.controls:
+                out.append([cpt.x, cpt.y])
+        points_n.append(out)
+
+    points_n = torch.tensor(points_n, requires_grad=True)
+    color = torch.tensor([s.color for s in squad], requires_grad=True)
+
+    paths = [s.as_path() for s in squad]
+    path_groups = [pydiffvg.ShapeGroup(shape_ids=torch.tensor([i]),
+                                       fill_color=torch.tensor(squad[i].color))
+                   for i in range(len(squad))]
+    scene_args = pydiffvg.RenderFunction.serialize_scene(
+        256, 256, paths, path_groups
+    )
+    img = render(256,  # width
+                 256,  # height
+                 2,   # num_samples_x
+                 2,   # num_samples_y
+                 1,   # seed
+                 None,
+                 *scene_args)
+
+    points, controls, color = [torch.tensor(x, requires_grad=True)
+                               for x in gm.to_numbers()]
+
+    optimizer = torch.optim.Adam([points, color, points_n], lr=1e-2)
+
+    for t in range(180):
+        print(f"iteration {t}")
+        optimizer.zero_grad()
+
+        points_n.data = torch.tensor(
+            GradientMesh.from_path_points(points_n, color).to_path_points()
+        )
+
+        for i in range(len(paths)):
+            paths[i].points = points_n[i] * 256
+
+        for i in range(len(path_groups)):
+            path_groups[i].fill_color = color[i]
+
+        scene_args = pydiffvg.RenderFunction.serialize_scene(
+            256, 256, paths, path_groups)
+
+        img = render(256,   # width
+                     256,   # height
+                     2,     # num_samples_x
+                     2,     # num_samples_y
+                     t+1,   # seed
+                     None,
+                     *scene_args)
+
+        pydiffvg.imwrite(img.cpu(),
+                         f'test_data/test_curve/iter_{filename}_'
+                         f'{str(t).zfill(5)}.png',
+                         gamma=2.2)
+
+        loss = (img - target).pow(2).sum()
+
+        loss.backward()
+
+        print(f'loss: {loss}')
+        print(f'points.grad {points.grad}')
+        print(f'color.grad {color.grad}')
+
+        optimizer.step()
+
+
+def slideshow(n=30, s=1, do_mask=False):
+    mask = None
+    for i in range(n):
+        if do_mask:
+            mask = [1] * 4
+            print(i % n)
+            mask[i % 4] = 0
+            print(mask)
+
+        mult_quad_test(mask=mask)
+        sleep(s)
+
+
+def get_mesh():
+    a, b, c, d = quads()
+    join_quads(a,b,c,d)
+
+    gm = GradientMesh(a, b, c, d)
+    return gm