Major refactors

.gitignore

@@ -1,4 +1,5 @@
 build
+gradientmesh/test_data/examples
 apps/results
 apps/files
 apps/__pycache__

gradientmesh/gmtypes.py

@@ -4,36 +4,21 @@ import torch
 from random import uniform
 import math
 
-import pydiffvg
+from util import rgb2hex, any_map
 
-from util import rgb2hex
 
 class Point:
     """2D point, optionally with control points;
         alternatively, 2-vector"""
     def __init__(self, x: float, y: float,
-                 controls: list[Point] = None,
                  round=False):
 
-        if isinstance(x, torch.Tensor):
-            # Convert from tensor
-            x = x.item()
-            y = y.item()
-
-        self.id = id(self)
         self.x = x
         self.y = y
-        self.controls = controls or []
 
         if round:
             self.round()
 
-    def add_control(self, control: Point):
-        self.controls.append(control)
-
-    def add_to_patch(self, patch: Patch):
-        self.patches.append(patch)
-
     def as_xy(self):
         return [self.x, self.y]
 
@@ -41,27 +26,19 @@ class Point:
         self.x = int(self.x * 100) / 100.0
         self.y = int(self.y * 100) / 100.0
 
-    def replace(self, pt: Point):
+        return self
-        """"Replace (x,y) coordinates of point while maintaining pointer."""
-        self.x = pt.x
-        self.y = pt.y
 
     def add(self, pt: Point):
         self.x += pt.x
         self.y += pt.y
-        for cp in self.controls:
+
-            cp.x += pt.x
+        return self
-            cp.y += pt.y
 
     def mult(self, pt: Point):
         self.x *= pt.x
         self.y *= pt.y
-        for cp in self.controls:
-            cp.x *= pt.x
-            cp.y *= pt.y
 
-    def equalize(self, other):
+        return self
-        self.id = other.id
 
     @classmethod
     def random(cls, rx=(0, 1), ry=(0, 1)):
@@ -73,67 +50,48 @@ class Point:
         out.round()
         return out
 
-    def __eq__(self, other):
-        return self.id == other.id
-
     def __hash__(self):
-        return self.id
+        # Used for removing duplicate points
+        return id(self)
 
     def __repr__(self):
-        return f"P<({self.x}, {self.y})[{len(self.controls)}]>"
+        return f"P<({self.x}, {self.y})>"
 
     def __str__(self):
         return self.__repr__()
 
 
 class Patch:
-    """Cubic patch."""
+    """Bicubic patch."""
-    def __init__(self, points: list[Point], color=(0.2, 0.5, 0.7, 1.0)):
+    def __init__(self,
+                 points: list[Point],
+                 controls: list[list[Point]],
+                 color=(0.2, 0.5, 0.7, 1.0)):
         self.points = points
+        self.controls = controls
         self.color = color
 
     def translate(self, pt: Point):
         for p in self.points:
             p.add(pt)
 
+        for q in self.controls:
+            for p in q:
+                p.add(pt)
+
     def scale(self, pt: Point):
         for p in self.points:
             p.mult(pt)
 
-    def as_path(self, width=256, height=256) -> pydiffvg.Path:
+        for q in self.controls:
-        ppoints = []
+            for p in q:
-        for pt in self.points:
+                p.mult(pt)
-            ppoints.append([pt.x * width, pt.y * height])
-            for cpt in pt.controls:
-                ppoints.append([cpt.x * width, cpt.y * height])
-
-        return pydiffvg.Path(
-            num_control_points=torch.tensor(
-                [len(p.controls) for p in self.points]
-            ),
-            points=torch.tensor(ppoints, requires_grad=True),
-            is_closed=True
-        )
-
-    def as_shape_group(self, color=None) -> pydiffvg.ShapeGroup:
-        # TODO proper id handling
-        return pydiffvg.ShapeGroup(
-            shape_ids=torch.tensor([0]),
-            fill_color=torch.tensor(color or self.color, requires_grad=True)
-        )
-
-    def get_points(self):
-        out = []
-        for p in self.points:
-            out.append([p.x, p.y])
-            for cp in p.controls:
-                out.append([cp.x, cp.y])
-
-        return out
 
     @classmethod
     def random(cls, degree=4, num_control_points=2):
         num_control_points = [num_control_points] * degree
+        """Returns a random Patch with `degree` vertices
+        and `num_control_points` control points per edge."""
 
         # Random tweaks to regular polygon base
         angle = 2 * math.pi / degree
@@ -156,6 +114,7 @@ class Patch:
 
             points.append(pt)
 
+        control_points = []
         for i in range(len(num_control_points)):
             pt = points[i]
             npt = points[i+1 if i+1 < degree else 0]
@@ -164,15 +123,14 @@ class Patch:
             dx = (npt.x - pt.x) / (ncp + 1)
             dy = (npt.y - pt.y) / (ncp + 1)
 
-            for j in range(1, ncp+1):
+            control_points.append([
-                midpoint = Point(
+                Point(
                     pt.x + j * dx * uniform(0.8, 1.2) + uniform(0, 0.2),
                     pt.y + j * dy * uniform(0.8, 1.2) + uniform(0, 0.2)
-                )
+                ) for j in range(1, ncp+1)
+            ])
 
-                pt.add_control(midpoint)
+        out = cls(points, control_points)
-
-        out = cls(points)
         out.color = (
             uniform(0, 1),
             uniform(0, 1),
@@ -189,8 +147,11 @@ class Patch:
 
 
 class Quad(Patch):
+    """Quadrilateral bicubic patch."""
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
+        # Assign 4 vertices as self.{top, bototm, left, right}
         self.top, self.right, self.bottom, self.left = self.points
         by_y = sorted(self.points, key=lambda pt: pt.y)
         self.top, self.bottom = by_y[0], by_y[-1]
@@ -199,184 +160,141 @@ class Quad(Patch):
                       key=lambda pt: pt.x)
         self.left, self.right = by_x
 
-        for pt in by_y:
+        # Assign 4 edges as self.{northeat, southeast, southwest, northwest}
-            pt.round()
+        self.northeast = self.controls[self.points.index(self.top)]
+        self.southeast = self.controls[self.points.index(self.right)]
+        self.southwest = self.controls[self.points.index(self.bottom)]
+        self.northwest = self.controls[self.points.index(self.left)]
 
-    @classmethod
+        self.set_points()
-    def from_path_points(cls, pp, color=None):
-        pt = None
-        points = []
-        for i in range(len(pp)):
-            if i % (len(pp) // 4) == 0:
-                if pt:
-                    points.append(pt)
-                pt = Point(*pp[i])
-            else:
-                pt.add_control(Point(*pp[i]))
 
-        points.append(pt)
+    def set_points(self):
-        return cls(points, color)
+        """Reset self.points and self.controls; used after mutating
+        e.g. self.top, self.southwest &c."""
+        self.points = [self.left, self.top,
+                       self.right, self.bottom]
 
-    def to_path_points(self):
+        self.controls = [self.northwest, self.northeast,
+                         self.southeast, self.southwest]
+
+
+class PointMapping:
+    """Mapping of unique points in a mesh to separate shapes for diffvg."""
+    def __init__(self, points, controls, raw_points, colors):
+        self.points = points
+        self.controls = controls
+        self.raw_points = raw_points
+        self.data = torch.tensor(
+            [pt.as_xy() for pt in raw_points], requires_grad=True
+        )
+        self.colors = colors
+
+    def as_shapes(self):
         out = []
-        for pt in self.points:
-            out.append(pt.as_xy())
-            for cp in pt.controls:
-                out.append(cp.as_xy())
 
-        return out
+        for i in range(len(self.points)):
+            quad = self.points[i]
+
+            quadpoints = []
+            for j in range(len(quad)):
+                quadpoints.append(quad[j])
+                quadpoints += self.controls[i][j]
+
+            out.append(quadpoints)
+
+        return [torch.stack(x) for x in
+                any_map(lambda idx: self.data[idx], out)]
 
 
 class GradientMesh:
+    """Bicubic quadrilateral mesh."""
     def __init__(self, *quads: Quad):
         self.quads = quads
 
-    def as_shape_groups(self):
+    def as_mapping(self):
-        sg = [quad.as_shape_group() for quad in self.quads]
+        """Convert GradientMesh to PointMapping"""
-        for i in range(len(sg)):
-            sg[i].shape_ids = torch.tensor([i])
-        return sg
-
-    def as_shapes(self, width, height):
-        return [quad.as_path(width, height) for quad in self.quads]
-
-    def to_numbers(self):
         points = []
         controls = []
-        for quad in self.quads:
+        raw_points = []
-            qp = []
-            qcp = []
-            for p in quad.points:
-                qp.append([p.x, p.y])
-                qcp.append([[cp.x, cp.y] for cp in p.controls])
-            points.append(qp)
-            controls.append(qcp)
 
-        return [points, controls, [q.color for q in self.quads]]
-
-    def from_numbers(self, numbers):
-        points, controls, colors = numbers
-        for q in range(4):
-            self.quads[q].color = colors[q]
-            pts = points[q]
-            ctrls = controls[q]
-            for p in range(4):
-                self.quads[q].points[p].replace(Point(*pts[p]))
-                for c in range(len(self.quads[q].points[p].controls)):
-                    self.quads[q].points[p].controls[c].replace(
-                        Point(*ctrls[p][c])
-                    )
-
-    @classmethod
-    def from_path_points(cls, pp, colors):
-        a, b, c, d = [Quad.from_path_points(pp[x], colors[x])
-                      for x in range(len(pp))]
-        join_quads(a, b, c, d, scale=False, translate=False)
-
-        return cls(a, b, c, d)
-
-    def to_path_points(self):
-        out = []
         for q in self.quads:
-            out.append(q.to_path_points())
+            points.append(q.points)
-        return out
+            controls.append(q.controls)
+            raw_points += q.points
+            for cp in q.controls:
+                raw_points += cp
 
-    def to_point_map(self):
+        raw_points = list(set(raw_points))
-        # XXX this doesn't work
-        # because of control points
-        pts: list[Point] = []
-        template: list[list[int]] = []
 
-        for quad in self.quads:
+        points = any_map(lambda p: raw_points.index(p), points)
-            for pt in quad.points:
+        controls = any_map(lambda p: raw_points.index(p), controls)
-                pts.append(pt)
 
-        pts = list(dict.fromkeys(pts))
+        colors = torch.tensor([q.color for q in self.quads],
+                              requires_grad=True)
 
-        for idx in range(len(self.quads)):
+        return PointMapping(points, controls, raw_points, colors)
-            template.append([
-                pts.index(pt) for pt in self.quads[idx].points
-            ])
 
-        return (pts, template)
+    def from_mapping(self, mapping: PointMapping):
-
+        pass
-    def from_point_map(self, pts, template):
-        # XXX this not either
-        mapped = []
-        for q in template:
-            mapped.append([
-                pts[x] for x in q
-            ])
-
-        return mapped
 
 
-def average_points(points: list[Point]) -> Point:
+def average_points(*points: list[Point]) -> Point:
+    """Average (i.e. geometric center) of points."""
     x = sum([pt.x for pt in points]) / len(points)
     y = sum([pt.y for pt in points]) / len(points)
     return Point(x, y)
 
 
-def equalize_points(points: list[Point]):
+def join_quads(a: Quad, b: Quad, c: Quad, d: Quad,
-    first = points[0]
+               scale=True, translate=True,
-    for pt in points:
+               step=10):
-        pt.equalize(first)
+    """Join 4 quadrilaterals so that they form a mesh with a center point."""
 
+    def combine_cp(cpa, cpb):
+        return [average_points(a, b) for a, b in zip(cpa, cpb)]
 
-def equalize_cp(points: list[Point]):
+    # If each quad occupies full space, make it so that they occupy 1/4 of space
-    first = points[0].controls[0]
-    for pt in points:
-        for cp in pt.controls:
-            cp.equalize(first)
-
-
-def merge_points(points: list[Point]):
-    merged = average_points(points)
-    for pt in points:
-        pt.replace(merged)
-
-
-def merge_cp(points: list[Point]):
-    ct1 = points[0].controls
-    ct2 = points[1].controls[::-1]
-    for i in range(len(ct1)):
-        merged = average_points([ct1[i], ct2[i]])
-        ct1[i].replace(merged)
-        ct2[i].replace(merged)
-
-
-def join_quads(a: Quad, b: Quad, c: Quad, d: Quad, scale=True, translate=True):
-    if translate:
-        b.translate(a.top)
-        c.translate(a.right)
-        d.translate(a.bottom)
-
-    merge_points([a.right, b.bottom, c.left, d.top])
-
-    merge_points([a.top, b.left])
-    merge_points([a.bottom, d.left])
-    merge_points([b.right, c.top])
-    merge_points([c.bottom, d.right])
-
     if scale:
         a.scale(Point(0.5, 0.5))
         b.scale(Point(0.5, 0.5))
         c.scale(Point(0.5, 0.5))
         d.scale(Point(0.5, 0.5))
 
-    merge_cp([a.right, d.left])
+    # If quads are on top of each other, translate so they are not
-    merge_cp([a.top, b.bottom])
+    if translate:
-    merge_cp([c.left, b.right])
+        b.translate(a.top)
-    merge_cp([c.bottom, d.top])
+        c.translate(a.right)
+        d.translate(a.bottom)
 
-    equalize_points([a.right, b.bottom, c.left, d.top])
+    # Equalize centerpoint
-    equalize_points([a.top, b.left])
+    a.right, b.bottom, c.left, d.top = [
-    equalize_points([a.bottom, d.left])
+        average_points(a.right, b.bottom, c.left, d.top)
-    equalize_points([b.right, c.top])
+    ] * 4
-    equalize_points([c.bottom, d.right])
 
-    equalize_cp([a.right, d.left])
+    # Equalize non-center shared points
-    equalize_cp([a.top, b.bottom])
+    a.top, b.left = [average_points(a.top, b.left)] * 2
-    equalize_cp([c.left, b.right])
+    a.bottom, d.left = [average_points(a.bottom, d.left)] * 2
-    equalize_cp([c.bottom, d.top])
+    b.right, c.top = [average_points(b.right, c.top)] * 2
+    c.bottom, d.right = [average_points(c.bottom, d.right)] * 2
 
+    # Equalize edges
+    a.northeast, b.southwest = (
+        (comb := combine_cp(a.northeast, b.southwest)),
+        comb[::-1]
+    )
+    b.southeast, c.northwest = (
+        (comb := combine_cp(b.southeast, c.northwest)),
+        comb[::-1]
+    )
+    c.southwest, d.northeast = (
+        (comb := combine_cp(c.southwest, d.northeast)),
+        comb[::-1]
+    )
+    d.northwest, a.southeast = (
+        (comb := combine_cp(d.northwest, a.southeast)),
+        comb[::-1]
+    )
+
+    # Update points
+    for q in [a, b, c, d]:
+        q.set_points()

gradientmesh/test.py

@@ -7,171 +7,95 @@ import torch
 import random
 from random import uniform
 
-from gmtypes import GradientMesh, Quad, Patch, Point, join_quads
+from gmtypes import GradientMesh, Quad, PointMapping, join_quads
+
+def get_mesh() -> GradientMesh:
+    """Helper function to get a random mesh."""
+    a, b, c, d = [Quad.random() for _ in range(4)]
+    join_quads(a, b, c, d)
+    return GradientMesh(a, b, c, d)
 
 
-def quads():
+def render_mesh(mesh: PointMapping,
-    return [
+                filename='test_data/mesh.png',
-        Quad.random(),
+                width=1024,
-        Quad.random(),
+                height=1024,
-        Quad.random(),
+                num_control_points=2,
-        Quad.random(),
+                seed=None):
-    ]
 
+    random.seed(seed)
 
-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
+    ncp = torch.tensor([num_control_points] * len(mesh.points))
 
-    patch = Patch.random()
+    # Scale
+    # TODO non-uniform scaling
+    points = [x * width for x in mesh.as_shapes()]
 
-    shape_groups = [patch.as_shape_group()]
+    shapes = [
-    shapes = [patch.as_path(width, height)]
+        pydiffvg.Path(num_control_points=ncp,
+                      points=pts,
+                      is_closed=True)
+        for pts in points
+    ]
 
-    scene_args = pydiffvg.RenderFunction.serialize_scene(width, height,
+    shape_groups = [
-                                                         shapes, shape_groups)
+        pydiffvg.ShapeGroup(shape_ids=torch.tensor([i]),
+                            fill_color=mesh.colors[i])
+        for i in range(len(mesh.points))
+    ]
 
-    img = render(width, height, 2, 2, 0, None, *scene_args)
+    scene_args = pydiffvg.RenderFunction.serialize_scene(
-    pydiffvg.imwrite(img.cpu(), f"test_data/{filename}", gamma=2.2)
+        width,
+        height,
+        shapes,
+        shape_groups
+    )
+
+    img = render(width,
+                 height,
+                 2,  # num smaples x
+                 2,  # num samples y
+                 0,  # seed
+                 None,
+                 *scene_args)
+    pydiffvg.imwrite(img.cpu(), filename, gamma=2.2)
 
     return img
 
 
-def mult_quad_test(filename='multiple_quads.png', width=1024,
+def test_render(filename='test_data/target.png',
-                   height=1024, num_control_points=None, mask=None, seed=None):
+                width=1024,
-    random.seed(seed)
+                height=1024):
-    mask = mask or [1, 1, 1, 1]
+    return render_mesh(get_mesh().as_mapping(),
-    pydiffvg.set_use_gpu(torch.cuda.is_available())
+                       width=width,
-    render = pydiffvg.RenderFunction.apply
+                       height=height,
-
+                       filename=filename)
-    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():
+def optimize():
-    filename = 'optimize_test.png'
+    width, height = 256, 256
-    pydiffvg.set_use_gpu(torch.cuda.is_available())
+    target = test_render(width=width, height=height).clone()
-    render = pydiffvg.RenderFunction.apply
 
-    target = mult_quad_test(width=256, height=256)
+    mesh = get_mesh().as_mapping()
 
-    squad = quads()
+    optimizer = torch.optim.Adam([mesh.data, mesh.colors], lr=1e-2)
 
-    join_quads(*squad)
+    for t in range(150):
-
-    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(
+        img = render_mesh(mesh,
-            GradientMesh.from_path_points(points_n, color).to_path_points()
+                          filename=f"test_data/mesh_optim_{str(t).zfill(3)}.png",
-        )
+                          width=width,
-
+                          height=height)
-        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()
+        # FIXME no need to retain graph
+        loss.backward(retain_graph=True)
 
         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

gradientmesh/util.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env ipython
+from types import FunctionType
+
 def clamp(val, low, high):
     return min(high, max(low, val))
 
@@ -10,3 +12,14 @@ def rgb2hex(r, g, b, a=None):
         int(b * 255)
     )
     return hex_value
+
+
+def any_map(f: FunctionType, lst: list):
+    result = []
+    for itm in lst:
+        if isinstance(itm, list):
+            result.append(any_map(f, itm))
+        else:
+            result.append(f(itm))
+
+    return result