301 lines
8.5 KiB
Python
301 lines
8.5 KiB
Python
#!/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()
|