#!/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])