feat(badapple): polygon tracing and rasterisation

bad-apple/pipeline/vectorize.py

@@ -0,0 +1,170 @@
+"""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())
+
+    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)

bad-apple/tests/test_vectorize.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+from pipeline.vectorize import rdp, rasterize, trace
+
+
+def test_rdp_drops_collinear_point():
+    pts = [(0, 0), (1, 0), (2, 0)]
+    assert rdp(pts, 0.5) == [(0, 0), (2, 0)]
+
+
+def test_rasterize_square():
+    poly = [[(1, 1), (4, 1), (4, 4), (1, 4), (1, 1)]]
+    out = rasterize([poly], 6, 6)
+    assert out[2, 2]
+    assert not out[0, 0]
+
+
+def test_rasterize_square_with_hole():
+    outer = [(1, 1), (5, 1), (5, 5), (1, 5), (1, 1)]
+    hole = [(2, 2), (2, 4), (4, 4), (4, 2), (2, 2)]
+    out = rasterize([[outer, hole]], 7, 7)
+    assert out[3, 1]
+    assert not out[3, 3]
+
+
+def test_trace_simple_shape_and_roundtrip():
+    mask = np.zeros((16, 16), dtype=bool)
+    mask[4:12, 4:12] = True
+    polys = trace(mask, {"rdp_epsilon": 1.0, "min_segment": 2})
+    assert len(polys) >= 1
+    roundtrip = rasterize(polys, 16, 16)
+    assert roundtrip[6, 6]
+    assert not roundtrip[0, 0]