diff --git a/.gitignore b/.gitignore index bb2ce51..cf10f5c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ build +gradientmesh/test_data/examples apps/results apps/files apps/__pycache__ diff --git a/gradientmesh/gmtypes.py b/gradientmesh/gmtypes.py index c8c97be..20543fb 100644 --- a/gradientmesh/gmtypes.py +++ b/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): - """"Replace (x,y) coordinates of point while maintaining pointer.""" - self.x = pt.x - self.y = pt.y + return self 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 + + return self 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 + return self @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.""" - def __init__(self, points: list[Point], color=(0.2, 0.5, 0.7, 1.0)): + """Bicubic patch.""" + 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: - 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 + for q in self.controls: + for p in q: + p.mult(pt) @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): - midpoint = Point( + control_points.append([ + 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) + out = cls(points, control_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: - pt.round() + # Assign 4 edges as self.{northeat, southeast, southwest, northwest} + 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 - 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])) + self.set_points() - points.append(pt) - return cls(points, color) + def set_points(self): + """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): - 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): + def as_mapping(self): + """Convert GradientMesh to PointMapping""" 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) + raw_points = [] - 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 + points.append(q.points) + controls.append(q.controls) + raw_points += q.points + for cp in q.controls: + raw_points += cp - def to_point_map(self): - # XXX this doesn't work - # because of control points - pts: list[Point] = [] - template: list[list[int]] = [] + raw_points = list(set(raw_points)) - for quad in self.quads: - for pt in quad.points: - pts.append(pt) + points = any_map(lambda p: raw_points.index(p), points) + controls = any_map(lambda p: raw_points.index(p), controls) - 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)): - template.append([ - pts.index(pt) for pt in self.quads[idx].points - ]) + return PointMapping(points, controls, raw_points, colors) - 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 from_mapping(self, mapping: PointMapping): + pass -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]): - first = points[0] - for pt in points: - pt.equalize(first) +def join_quads(a: Quad, b: Quad, c: Quad, d: Quad, + scale=True, translate=True, + step=10): + """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]): - 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 each quad occupies full space, make it so that they occupy 1/4 of space 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]) + # If quads are on top of each other, translate so they are not + if translate: + b.translate(a.top) + c.translate(a.right) + d.translate(a.bottom) - 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 centerpoint + a.right, b.bottom, c.left, d.top = [ + average_points(a.right, b.bottom, c.left, d.top) + ] * 4 - equalize_cp([a.right, d.left]) - equalize_cp([a.top, b.bottom]) - equalize_cp([c.left, b.right]) - equalize_cp([c.bottom, d.top]) + # Equalize non-center shared points + a.top, b.left = [average_points(a.top, b.left)] * 2 + a.bottom, d.left = [average_points(a.bottom, d.left)] * 2 + 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() diff --git a/gradientmesh/test.py b/gradientmesh/test.py index 5f15de8..a6f42c5 100644 --- a/gradientmesh/test.py +++ b/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(): - return [ - Quad.random(), - Quad.random(), - Quad.random(), - Quad.random(), - ] +def render_mesh(mesh: PointMapping, + filename='test_data/mesh.png', + width=1024, + height=1024, + num_control_points=2, + 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 = [patch.as_path(width, height)] + shapes = [ + pydiffvg.Path(num_control_points=ncp, + points=pts, + is_closed=True) + for pts in points + ] - scene_args = pydiffvg.RenderFunction.serialize_scene(width, height, - shapes, shape_groups) + 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) - pydiffvg.imwrite(img.cpu(), f"test_data/{filename}", gamma=2.2) + scene_args = pydiffvg.RenderFunction.serialize_scene( + 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, - 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 test_render(filename='test_data/target.png', + width=1024, + height=1024): + return render_mesh(get_mesh().as_mapping(), + width=width, + height=height, + filename=filename) -def om(): - filename = 'optimize_test.png' - pydiffvg.set_use_gpu(torch.cuda.is_available()) - render = pydiffvg.RenderFunction.apply +def optimize(): + width, height = 256, 256 + target = test_render(width=width, height=height).clone() - 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) - - 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): + for t in range(150): 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) + img = render_mesh(mesh, + filename=f"test_data/mesh_optim_{str(t).zfill(3)}.png", + width=width, + height=height) 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 diff --git a/gradientmesh/test_data/meme.png b/gradientmesh/test_data/meme.png deleted file mode 100644 index 9c41291..0000000 Binary files a/gradientmesh/test_data/meme.png and /dev/null differ diff --git a/gradientmesh/test_data/monotone_quad.png b/gradientmesh/test_data/monotone_quad.png deleted file mode 100644 index e247ba6..0000000 Binary files a/gradientmesh/test_data/monotone_quad.png and /dev/null differ diff --git a/gradientmesh/test_data/multiple_quads.png b/gradientmesh/test_data/multiple_quads.png deleted file mode 100644 index ebd964e..0000000 Binary files a/gradientmesh/test_data/multiple_quads.png and /dev/null differ diff --git a/gradientmesh/test_data/optimize_test.png b/gradientmesh/test_data/optimize_test.png deleted file mode 100644 index a331f6b..0000000 Binary files a/gradientmesh/test_data/optimize_test.png and /dev/null differ diff --git a/gradientmesh/test_data/optimize_test.png.png b/gradientmesh/test_data/optimize_test.png.png deleted file mode 100644 index 1bda8eb..0000000 Binary files a/gradientmesh/test_data/optimize_test.png.png and /dev/null differ diff --git a/gradientmesh/test_data/random_quad.png b/gradientmesh/test_data/random_quad.png deleted file mode 100644 index 30ba6c7..0000000 Binary files a/gradientmesh/test_data/random_quad.png and /dev/null differ diff --git a/gradientmesh/util.py b/gradientmesh/util.py index 842a379..21860dd 100644 --- a/gradientmesh/util.py +++ b/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