"""Turn a binary frame into simplified straight-edged polygons, and back.""" from __future__ import annotations import json import os import subprocess import tempfile from typing import Iterable, Sequence import numpy as np Point = tuple[float, float] Ring = list[Point] Polygon = list[Ring] def _distance_to_line(p: Point, a: Point, b: Point) -> float: if a == b: return ((p[0] - a[0]) ** 2 + (p[1] - a[1]) ** 2) ** 0.5 ax, ay = a bx, by = b px, py = p dx, dy = bx - ax, by - ay return abs(dy * px - dx * py + bx * ay - by * ax) / ((dx * dx + dy * dy) ** 0.5) def rdp(points: Sequence[Point], epsilon: float) -> list[Point]: """Ramer-Douglas-Peucker on an open polyline.""" pts = list(points) if len(pts) <= 2: return pts a, b = pts[0], pts[-1] worst = 0.0 index = 0 for i in range(1, len(pts) - 1): d = _distance_to_line(pts[i], a, b) if d > worst: worst, index = d, i if worst <= epsilon: return [a, b] left = rdp(pts[:index + 1], epsilon) right = rdp(pts[index:], epsilon) return left[:-1] + right def rdp_closed(points: Sequence[Point], epsilon: float) -> list[Point]: """Simplify a closed ring, keeping it closed.""" pts = list(points) if len(pts) > 1 and pts[0] == pts[-1]: pts = pts[:-1] if len(pts) < 3: return list(points) # Split at the point farthest from the first so the endpoints are not # coincident (RDP degenerates when a == b). far = max(range(1, len(pts)), key=lambda i: (pts[i][0] - pts[0][0]) ** 2 + (pts[i][1] - pts[0][1]) ** 2) first = rdp(pts[:far + 1], epsilon) second = rdp(pts[far:] + [pts[0]], epsilon) ring = first[:-1] + second if ring[0] != ring[-1]: ring.append(ring[0]) return ring def parse_potrace_geojson(text: str) -> list[Polygon]: """Parse potrace's GeoJSON into polygons of (outer, holes...) rings.""" data = json.loads(text) out: list[Polygon] = [] for feature in data.get("features", []): geometry = feature.get("geometry") or {} kind = geometry.get("type") coordinates = geometry.get("coordinates") or [] if kind == "Polygon": polygons = [coordinates] elif kind == "MultiPolygon": polygons = coordinates else: continue for rings in polygons: poly: Polygon = [] for ring in rings: poly.append([(float(x), float(y)) for x, y in ring]) if poly: out.append(poly) return out def trace(mask: np.ndarray, params: dict) -> list[Polygon]: """Trace a black mask into simplified straight-line polygons.""" if mask.ndim != 2: raise ValueError("trace wants a 2-D mask") height, width = mask.shape if width % 8: raise ValueError("potrace input width must be a multiple of 8") epsilon = float(params.get("rdp_epsilon", 1.5)) min_segment = float(params.get("min_segment", 2)) pbm = np.packbits(mask.astype(bool), axis=1, bitorder="big").tobytes() with tempfile.TemporaryDirectory(prefix="ba-trace-") as tmp: pbm_path = os.path.join(tmp, "frame.pbm") json_path = os.path.join(tmp, "frame.json") with open(pbm_path, "wb") as handle: handle.write(b"P4\n%d %d\n" % (width, height)) handle.write(pbm) subprocess.run( ["potrace", "-b", "geojson", "-a", "0", "-t", "4", "-u", "1", "-o", json_path, pbm_path], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) with open(json_path, "r", encoding="utf-8") as handle: polygons = parse_potrace_geojson(handle.read()) # potrace's GeoJSON uses a y-up coordinate system; the framebuffer is # y-down. Without this the whole movie is vertically mirrored. flipped: list[Polygon] = [] for poly in polygons: rings: Polygon = [] for ring in poly: rings.append([(x, float(height) - y) for x, y in ring]) flipped.append(rings) polygons = flipped simplified: list[Polygon] = [] for poly in polygons: rings: Polygon = [] for ring in poly: ring = rdp_closed(ring, epsilon) if len(ring) < 4: continue longest = max( ((ring[i][0] - ring[i - 1][0]) ** 2 + (ring[i][1] - ring[i - 1][1]) ** 2) ** 0.5 for i in range(1, len(ring)) ) if longest < min_segment: continue rings.append(ring) if rings: simplified.append(rings) return simplified def rasterize(polygons: Iterable[Polygon], width: int, height: int) -> np.ndarray: """Even-odd scanline fill of polygons (outer rings and holes).""" out = np.zeros((height, width), dtype=bool) edges: list[tuple[float, float, float, float]] = [] for poly in polygons: for ring in poly: for i in range(1, len(ring)): x0, y0 = ring[i - 1] x1, y1 = ring[i] if y0 == y1: continue edges.append((x0, y0, x1, y1)) if not edges: return out for y in range(height): scan = y + 0.5 xs: list[float] = [] for x0, y0, x1, y1 in edges: if (y0 <= scan < y1) or (y1 <= scan < y0): xs.append(x0 + (scan - y0) * (x1 - x0) / (y1 - y0)) xs.sort() for i in range(0, len(xs) - 1, 2): left, right = xs[i], xs[i + 1] start = max(0, int(np.ceil(left - 0.5))) end = min(width - 1, int(np.floor(right - 0.5))) if end >= start: out[y, start:end + 1] = True return out def vectorize_frame(mask: np.ndarray, params: dict) -> np.ndarray: """Trace and re-rasterise one logical frame.""" height, width = mask.shape return rasterize(trace(mask, params), width, height)