vt1200/bad-apple/pipeline/vectorize.py

"""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)