kerf-engine/engine/pipeline/vectorize.py

"""Vectorization pipeline — converts preprocessed binary images to SVG."""

import cv2
import numpy as np
import potrace
import vtracer


_TURNPOLICY_MAP = {
    "black": potrace.TURNPOLICY_BLACK,
    "white": potrace.TURNPOLICY_WHITE,
    "left": potrace.TURNPOLICY_LEFT,
    "right": potrace.TURNPOLICY_RIGHT,
    "minority": potrace.TURNPOLICY_MINORITY,
    "majority": potrace.TURNPOLICY_MAJORITY,
}


def potrace_trace(
    binary_img: np.ndarray,
    turdsize: int = 2,
    alphamax: float = 1.0,
    opticurve: bool = True,
    opttolerance: float = 0.2,
    turnpolicy: str = "minority",
) -> str:
    """Trace a binary image using Potrace and return an SVG string.

    Args:
        binary_img: 2D numpy array — nonzero pixels are foreground.
        turdsize: Despeckle threshold; curves with enclosed area below this are removed.
        alphamax: Corner detection threshold (0.0 = polygon, 1.3333 = no corners).
        opticurve: Whether to optimize curves by reducing Bezier segments.
        opttolerance: Tolerance for curve optimization.
        turnpolicy: How to resolve ambiguities — 'minority', 'majority', 'black',
                     'white', 'left', 'right'.

    Returns:
        Well-formed SVG string.
    """
    if binary_img.ndim != 2:
        raise ValueError(f"Expected 2D binary image, got shape {binary_img.shape}")

    h, w = binary_img.shape

    # Potrace interprets nonzero pixels as foreground.
    # Convert to uint32 — pypotrace needs values that fit in a C int.
    data = (binary_img > 0).astype(np.uint32)

    tp = _TURNPOLICY_MAP.get(turnpolicy, potrace.TURNPOLICY_MINORITY)

    bmp = potrace.Bitmap(data)
    path = bmp.trace(
        turdsize=turdsize,
        alphamax=alphamax,
        opticurve=int(opticurve),
        opttolerance=opttolerance,
        turnpolicy=tp,
    )

    return _path_to_svg(path, w, h)


def vtracer_trace(
    img: np.ndarray,
    colormode: str = "binary",
    hierarchical: str = "stacked",
    filter_speckle: int = 4,
    color_precision: int = 6,
    layer_difference: int = 16,
    corner_threshold: int = 60,
    length_threshold: float = 4.0,
    splice_threshold: int = 45,
    mode: str = "spline",
    path_precision: int | None = None,
    max_iterations: int = 10,
) -> str:
    """Trace an image using VTracer and return an SVG string.

    Unlike potrace_trace, this accepts both grayscale and color images.
    Internally encodes the image as PNG and passes it to VTracer's Rust backend.

    Args:
        img: 2D (grayscale) or 3D (BGR/BGRA) numpy array.
        colormode: 'color' or 'binary'.
        hierarchical: 'stacked' or 'cutout'.
        filter_speckle: Remove patches smaller than this area (in px).
        color_precision: Number of significant bits for color quantization (1-8).
        layer_difference: Delta threshold for color layer grouping.
        corner_threshold: Angle (degrees) below which a point is a corner.
        length_threshold: Minimum segment length before simplification.
        splice_threshold: Angle (degrees) for splicing splines.
        mode: 'spline', 'polygon', or 'none' — curve fitting strategy.
        path_precision: Decimal precision for path coordinates.
        max_iterations: Max curve-fitting iterations.

    Returns:
        SVG string (includes XML declaration and generator comment).
    """
    if img.ndim not in (2, 3):
        raise ValueError(f"Expected 2D or 3D image, got {img.ndim}D (shape {img.shape})")

    # Encode as PNG — VTracer accepts raw image bytes via convert_raw_image_to_svg.
    ok, buf = cv2.imencode(".png", img)
    if not ok:
        raise RuntimeError("Failed to encode image as PNG for VTracer")

    # Build kwargs, omitting None values so VTracer uses its defaults.
    kwargs: dict = dict(
        img_format="png",
        colormode=colormode,
        hierarchical=hierarchical,
        filter_speckle=filter_speckle,
        color_precision=color_precision,
        layer_difference=layer_difference,
        corner_threshold=corner_threshold,
        length_threshold=length_threshold,
        splice_threshold=splice_threshold,
        mode=mode,
        max_iterations=max_iterations,
    )
    if path_precision is not None:
        kwargs["path_precision"] = path_precision

    return vtracer.convert_raw_image_to_svg(buf.tobytes(), **kwargs)


def _path_to_svg(path, width: int, height: int) -> str:
    """Convert a potrace Path object to an SVG string."""
    parts = []
    for curve in path:
        sx, sy = curve.start_point
        parts.append(f"M {sx:.3f},{sy:.3f}")
        for segment in curve.segments:
            if segment.is_corner:
                cx, cy = segment.c
                ex, ey = segment.end_point
                parts.append(f"L {cx:.3f},{cy:.3f} L {ex:.3f},{ey:.3f}")
            else:
                c1x, c1y = segment.c1
                c2x, c2y = segment.c2
                ex, ey = segment.end_point
                parts.append(
                    f"C {c1x:.3f},{c1y:.3f} {c2x:.3f},{c2y:.3f} {ex:.3f},{ey:.3f}"
                )
        parts.append("Z")

    d = " ".join(parts)

    return (
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{width}" height="{height}" '
        f'viewBox="0 0 {width} {height}">'
        f'<path d="{d}" fill="black" fill-rule="evenodd" stroke="none"/>'
        f"</svg>"
    )