many changes

gradientmesh/gmtypes.py

@@ -0,0 +1,382 @@
+#!/usr/bin/env ipython
+from __future__ import annotations
+import torch
+from random import uniform
+import math
+
+import pydiffvg
+
+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]
+
+    def round(self):
+        self.x = int(self.x * 100) / 100.0
+        self.y = int(self.y * 100) / 100.0
+
+    def replace(self, pt: Point):
+        """"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
+            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):
+        self.id = other.id
+
+    @classmethod
+    def random(cls, rx=(0, 1), ry=(0, 1)):
+        # Clamp to (0, 1)
+        rx = (max(0, rx[0]), min(1, max(0, rx[1])))
+        ry = (max(0, ry[0]), min(1, max(0, ry[1])))
+
+        out = cls(uniform(*rx), uniform(*ry))
+        out.round()
+        return out
+
+    def __eq__(self, other):
+        return self.id == other.id
+
+    def __hash__(self):
+        return self.id
+
+    def __repr__(self):
+        return f"P<({self.x}, {self.y})[{len(self.controls)}]>"
+
+    def __str__(self):
+        return self.__repr__()
+
+
+class Patch:
+    """Cubic patch."""
+    def __init__(self, points: list[Point], color=(0.2, 0.5, 0.7, 1.0)):
+        self.points = points
+        self.color = color
+
+    def translate(self, pt: Point):
+        for p in self.points:
+            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:
+        ppoints = []
+        for pt in self.points:
+            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
+
+        # Random tweaks to regular polygon base
+        angle = 2 * math.pi / degree
+        angle = uniform(0.8 * angle, 1.2 * angle)
+        points = []
+        for i in range(degree):
+            pt = Point(
+                uniform(0.3, 0.7) + 0.5 * math.cos(i * angle),
+                uniform(0.3, 0.7) + 0.5 * math.sin(i * angle)
+            )
+
+            # Stochastically clamp to (0,1)
+            for c in ['x', 'y']:
+                if (v := getattr(pt, c)) > 1:
+                    diff = v - 1
+                    setattr(pt, c, v - uniform(diff, 2 * diff))
+                elif v < 0:
+                    diff = -v
+                    setattr(pt, c, v + uniform(diff, 2 * diff))
+
+            points.append(pt)
+
+        for i in range(len(num_control_points)):
+            pt = points[i]
+            npt = points[i+1 if i+1 < degree else 0]
+
+            ncp = num_control_points[i]
+            dx = (npt.x - pt.x) / (ncp + 1)
+            dy = (npt.y - pt.y) / (ncp + 1)
+
+            for j in range(1, ncp+1):
+                midpoint = 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)
+                )
+
+                pt.add_control(midpoint)
+
+        out = cls(points)
+        out.color = (
+            uniform(0, 1),
+            uniform(0, 1),
+            uniform(0, 1),
+            0.7
+        )
+
+        return out
+
+    def __repr__(self):
+        out = f"F<({rgb2hex(*self.color)})"
+        out += f"[{', '.join([str(x) for x in self.points])}]>"
+        return out
+
+
+class Quad(Patch):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        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]
+        by_x = sorted([pt for pt in self.points
+                       if pt not in [self.top, self.bottom]],
+                      key=lambda pt: pt.x)
+        self.left, self.right = by_x
+
+        for pt in by_y:
+            pt.round()
+
+    @classmethod
+    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)
+        return cls(points, color)
+
+    def to_path_points(self):
+        out = []
+        for pt in self.points:
+            out.append(pt.as_xy())
+            for cp in pt.controls:
+                out.append(cp.as_xy())
+
+        return out
+
+
+class GradientMesh:
+    def __init__(self, *quads: Quad):
+        self.quads = quads
+
+    def as_shape_groups(self):
+        sg = [quad.as_shape_group() for quad in self.quads]
+        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:
+            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())
+        return out
+
+    def to_point_map(self):
+        # XXX this doesn't work
+        # because of control points
+        pts: list[Point] = []
+        template: list[list[int]] = []
+
+        for quad in self.quads:
+            for pt in quad.points:
+                pts.append(pt)
+
+        pts = list(dict.fromkeys(pts))
+
+        for idx in range(len(self.quads)):
+            template.append([
+                pts.index(pt) for pt in self.quads[idx].points
+            ])
+
+        return (pts, template)
+
+    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:
+    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]):
+    first = points[0]
+    for pt in points:
+        pt.equalize(first)
+
+
+def equalize_cp(points: list[Point]):
+    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])
+    merge_cp([a.top, b.bottom])
+    merge_cp([c.left, b.right])
+    merge_cp([c.bottom, d.top])
+
+    equalize_points([a.right, b.bottom, c.left, d.top])
+    equalize_points([a.top, b.left])
+    equalize_points([a.bottom, d.left])
+    equalize_points([b.right, c.top])
+    equalize_points([c.bottom, d.right])
+
+    equalize_cp([a.right, d.left])
+    equalize_cp([a.top, b.bottom])
+    equalize_cp([c.left, b.right])
+    equalize_cp([c.bottom, d.top])
+

gradientmesh/gradientmesh.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env ipython
+import pydiffvg
+import torch
+import skimage
+import numpy as np

gradientmesh/imgs/Triangle.obj

@@ -0,0 +1,5 @@
+v 100 150
+v 42.3 50
+v 157.7 50
+
+f 1 2 3

gradientmesh/parser.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env ipython
+import unittest
+from unittest import TestCase
+from dataclasses import dataclass
+
+
+class TestIO(TestCase):
+    pass
+
+class TestRender(TestCase):
+    pass
+
+@dataclass
+class Vertex:
+    index: int = None
+
+
+@dataclass
+class VertexData:
+    index: int = None
+    val: int = None
+    x: float = None
+    y: float = None
+    outIndex: int = None
+    looseEdgesIndex: list[int] = None
+
+
+@dataclass
+class FaceData:
+    index: int = None
+    val: int = None
+    sideIndex: int = None
+    resolution: int = None
+    tex: list[float] = None
+
+@dataclass
+class HalfEdgeData:
+    index: int = None
+    targIndex: int = None
+    twinIndex: int = None
+    prevIndex: int = None
+    nextIndex: int = None
+    polyIndex: int = None
+    r: float = None
+    g: float = None
+    b: float = None
+    tx: float = None
+    ty: float = None
+
+    def __repr__(self):
+        return f"HE[{self.index}]"
+
+
+def from_HEMESH(filename):
+    """Read from HEMESH file"""
+    indices: list[list[int]] = [[]]
+    vertices: list[VertexData] = []
+    faces: list[FaceData] = []
+    edges: list[HalfEdgeData] = []
+    with open(filename, 'r') as f:
+        if (ft := f.readline().strip()) != 'HEMESH':
+            raise ValueError(f"Unknown file type '{ft}'.")
+
+        numVert, numFace, numEdge = [int(x) for x in f.readline().split(' ')]
+
+        # Parse vertices
+        while len(vertices) < numVert:
+            if not (toks := f.readline().split(' ')):
+                continue
+
+            vert = VertexData()
+            vert.index = int(toks[0])
+            vert.val = int(toks[1])
+            vert.x, vert.y = float(toks[2]), float(toks[3])
+            vert.outIndex = int(toks[4])
+
+            numOutEdges = int(toks[5])
+            vert.looseEdgesIndex = [int(x) for x in toks[6:] if x.strip()]
+
+            vertices.append(vert)
+
+        # Faces
+        while len(faces) < numFace:
+            if not (toks := f.readline().split(' ')):
+                continue
+
+            face = FaceData()
+            (face.index, face.val,
+                face.sideIndex, face.res) = [int(x) for x in toks[:4]]
+            face.tex = [float(x) for x in toks[4:]]
+
+            faces.append(face)
+
+        while len(edges) < numEdge:
+            if not (toks := f.readline().split(' ')):
+                continue
+
+            edge = HalfEdgeData()
+
+            edge.index = int(toks[0])
+            edge.targIndex = int(toks[1])
+            edge.twinIndex = int(toks[2])
+            edge.prevIndex = int(toks[3])
+            edge.nextIndex = int(toks[4])
+            edge.polyIndex = int(toks[5])
+            edge.r = float(toks[6])
+            edge.g = float(toks[7])
+            edge.b = float(toks[8])
+            edge.tx = float(toks[9])
+            edge.ty = float(toks[10])
+
+            edges.append(edge)
+
+        return (vertices, faces, edges)

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

gradientmesh/test_data/epic.hemesh

@@ -0,0 +1,15 @@
+HEMESH
+4 1 8
+0 2 -0.517196 0.619048 0 0 
+1 2 0.25 0.25 3 0 
+2 2 -0.25 -0.25 1 0 
+3 2 0.25 -0.25 2 0 
+0 4 3 1 0 0 0 0 0 0.333333 0 0 0.333333 0 0 0.666667
+0 2 4 1 3 0 1 0 0 0 -0.909196 
+1 3 5 2 0 0 0 0 1 0.166667 0 
+2 1 6 3 1 0 1 1 0 0 0.166667 
+3 0 7 0 2 0 0 1 0 -0.851343 0 
+4 0 0 7 5 -1 0 0 1 0 0.909196 
+5 2 1 4 6 -1 1 1 0 -0.166667 0 
+6 3 2 5 7 -1 0 1 0 0 -0.166667 
+7 1 3 6 4 -1 1 0 0 0.404762 1.16667 

gradientmesh/test_data/square_brgb.hemesh

@@ -0,0 +1,15 @@
+HEMESH
+4 1 8
+0 2 -0.25 0.25 8 0 
+1 2 0.25 0.25 11 0 
+2 2 -0.25 -0.25 9 0 
+3 2 0.25 -0.25 10 0 
+0 4 11 1 1 0 0 0 0 1 0 0 0 0 1 0
+8 2 4 9 11 0 1 0 0 0 -0.166667 
+9 3 5 10 8 0 0 0 1 0.166667 0 
+10 1 6 11 9 0 1 1 0 0 0.166667 
+11 0 7 8 10 0 0 1 0 -0.166667 0 
+4 0 8 7 5 -1 0 0 1 0 0.166667 
+5 2 9 4 6 -1 1 1 0 -0.166667 0 
+6 3 10 5 7 -1 0 1 0 0 -0.166667 
+7 1 11 6 4 -1 1 0 0 0.166667 0 

gradientmesh/test_data/square_rbbb.hemesh

@@ -0,0 +1,15 @@
+HEMESH
+4 1 8
+0 2 -0.25 0.25 8 0 
+1 2 0.25 0.25 11 0 
+2 2 -0.25 -0.25 9 0 
+3 2 0.25 -0.25 10 0 
+0 4 11 1 0 0 1 0 0 1 0.666667 0 0 0 0 1
+8 2 4 9 11 0 1 0 0 0 -0.166667 
+9 3 5 10 8 0 0 0 1 0.166667 0 
+10 1 6 11 9 0 1 1 0 0 0.166667 
+11 0 7 8 10 0 0 1 0 -0.166667 0 
+4 0 8 7 5 -1 0 0 1 0 0.166667 
+5 2 9 4 6 -1 1 1 0 -0.166667 0 
+6 3 10 5 7 -1 0 1 0 0 -0.166667 
+7 1 11 6 4 -1 1 0 0 0.166667 0