#!/usr/bin/env ipython from __future__ import annotations import torch from random import uniform import math from util import rgb2hex, any_map class Point: """2D point, optionally with control points; alternatively, 2-vector""" def __init__(self, x: float, y: float, round=False): self.x = x self.y = y if round: self.round() 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 return self def add(self, pt: Point): self.x += pt.x self.y += pt.y return self def mult(self, pt: Point): self.x *= pt.x self.y *= pt.y return self @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 __hash__(self): # Used for removing duplicate points return id(self) def __repr__(self): return f"P<({self.x}, {self.y})>" def __str__(self): return self.__repr__() class Patch: """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) 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 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) control_points = [] 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) 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) ]) out = cls(points, control_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): """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] 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 # 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)] self.set_points() 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] 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 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_mapping(self): """Convert GradientMesh to PointMapping""" points = [] controls = [] raw_points = [] for q in self.quads: points.append(q.points) controls.append(q.controls) raw_points += q.points for cp in q.controls: raw_points += cp raw_points = list(set(raw_points)) points = any_map(lambda p: raw_points.index(p), points) controls = any_map(lambda p: raw_points.index(p), controls) colors = torch.tensor([q.color for q in self.quads], requires_grad=True) return PointMapping(points, controls, raw_points, colors) def from_mapping(self, mapping: PointMapping): pass 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 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)] # 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)) # 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 centerpoint a.right, b.bottom, c.left, d.top = [ average_points(a.right, b.bottom, c.left, d.top) ] * 4 # 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()