Spaces:
Running
Running
"""
ICON VECTORIZER - ACCURATE CENTERLINE
====================================
vectorizer_core.py — slow-but-accurate 4-stage centerline vectorization pipeline.
Converts black-outline icons (JPG/PNG) into clean centerline SVG with semantic
data-name labels (e.g. eyepiece-rim, optical-tube, curved-arm).
Pipeline
--------
Stage 0 - Preprocessing : 3x/4x INTER_CUBIC upscale, median+gaussian denoise,
user threshold, morphological opening.
Stage 1 - Segmentation : connected ink components (+ contour-hierarchy split),
optional MobileSAM refinement.
Stage 2 - Primitives : circularity test -> perfect <circle> (minEnclosingCircle
+ RANSAC refine); convex quad -> inset <rect>.
Stage 3 - Centerline : skimage skeletonize (Lee, NOT Zhang-Suen) + EDT distance,
skeleton graph (endpoints/junctions/loops), spur prune,
endpoint snap, junction branch-direction measurement.
Stage 4 - Fitting : RANSAC circle/ellipse on closed loops, Kasa arc fit,
Schneider cubic Bezier fit (Graphics Gems, tol=0.6px),
optional CLIP semantic labeling, clean SVG output.
Everything geometric is deterministic and offline. Only the OPTIONAL semantic
labeling (CLIP) and segmentation refinement (MobileSAM) need torch + model
downloads; the tracer works fully without them.
Author: built for Twins Design Studio / allfa — accuracy >> speed.
"""
from __future__ import annotations
import math
import os
import time
import urllib.request
from dataclasses import dataclass, field, replace
from typing import Callable, Dict, List, Optional, Tuple
import numpy as np
import cv2
from PIL import Image
try:
from skimage.morphology import skeletonize
HAS_SKIMAGE = True
except ImportError: # pragma: no cover
HAS_SKIMAGE = False
try:
from scipy.ndimage import gaussian_filter1d, distance_transform_edt
HAS_SCIPY = True
except ImportError: # pragma: no cover
HAS_SCIPY = False
# ----------------------------------------------------------------------------
# Configuration
# ----------------------------------------------------------------------------
SMOOTH_PRESETS = {"Low": 1, "Medium": 4, "High": 10}
@dataclass
class FitParams:
"""Fitting parameters (accurate mode).
tol is in ORIGINAL image pixels (user-facing unit); the pipeline works
at `scale`x, so the working tolerance is tol * scale.
"""
tol: float = 0.6 # max deviation in px (original image size)
smooth_iters: int = 1 # light pre-smoothing passes (Low)
reduce: bool = False # merge near-collinear cubics (OFF for accuracy)
# Prefer a single cubic over a fragmented fit when the chain is "almost"
# cubic: accept a single cubic if its max deviation is within this factor
# of tol. Keeps smooth curves as ONE editable cubic instead of dozens of
# fragments; genuine S-curves still split (their single-cubic error is far
# larger). The actual deviation is always reported per path.
SINGLE_CUBIC_TOL_FACTOR = 4.0
# Pre-fit smoothing sigma (upscaled px). Removes rasterization staircase
# noise (high frequency) without distorting real curve shape (which bends
# over tens of px). Applied `smooth_iters` times per SMOOTH_PRESETS.
SMOOTH_SIGMA = 1.5
@dataclass
class VectorizerConfig:
threshold: int = 110
upscale: int = 3 # Quality menu: 1x/2x/3x/4x
smoothing: str = "Low" # Low / Medium / High
tolerance: float = 0.6 # px, upscaled space
detect_corners: bool = True
detect_junctions: bool = True
detect_circles: bool = True # circles AND arcs
reduce_anchors: bool = False
use_sam: bool = False # MobileSAM refinement (slow, optional)
use_clip: bool = True # CLIP semantic labels (optional)
clip_prompts: List[str] = field(default_factory=lambda: [
"eyepiece", "optical tube", "lens", "microscope arm", "pivot",
"stage", "base", "handle", "knob", "button", "icon outline",
])
min_part_area: int = 60 # px^2 in upscaled space
def fit_params(self) -> FitParams:
return FitParams(
tol=float(self.tolerance),
smooth_iters=SMOOTH_PRESETS.get(self.smoothing, 1),
reduce=bool(self.reduce_anchors),
)
@dataclass
class Part:
id: int
mask: np.ndarray # bool array, upscaled working space, ink=True
bbox: Tuple[int, int, int, int] # x, y, w, h (upscaled)
label: str = "" # CLIP base label, e.g. "eyepiece"
# progress callback type: fn(fraction: float, stage: str)
ProgressCb = Optional[Callable[[float, str], None]]
def _progress(cb: ProgressCb, frac: float, stage: str):
if cb:
try:
cb(max(0.0, min(1.0, frac)), stage)
except Exception:
pass
# ============================================================================
# STAGE 0 - Preprocessing
# ============================================================================
def preprocess(gray: np.ndarray, cfg: VectorizerConfig) -> Tuple[np.ndarray, int]:
"""Upscale, denoise, threshold, clean. Returns (ink_mask_bool, scale)."""
scale = int(cfg.upscale)
h, w = gray.shape[:2]
big = cv2.resize(gray, (w * scale, h * scale), interpolation=cv2.INTER_CUBIC)
big = cv2.medianBlur(big, 3) # kill screen moire
big = cv2.GaussianBlur(big, (0, 0), 0.6 * scale) # gentle AA
_, bw = cv2.threshold(big, int(cfg.threshold), 255, cv2.THRESH_BINARY_INV)
bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8))
return (bw > 0), scale
def load_grayscale(path: str) -> np.ndarray:
img = Image.open(path).convert("L")
return np.array(img)
# ============================================================================
# STAGE 1 - Semantic segmentation
# ============================================================================
def _split_component_hierarchy(comp: np.ndarray) -> List[np.ndarray]:
"""Split one ink component via contour hierarchy (outer contours = parts).
Keeps it conservative: only splits when an outer contour (other than the
largest) is reasonably big; tiny fragments are merged back into the
largest region so thin touching strokes don't shatter.
"""
cnts, hier = cv2.findContours(comp.astype(np.uint8), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_SIMPLE)
if not cnts or hier is None:
return [comp]
hier = hier[0]
outers = [i for i in range(len(cnts)) if hier[i][3] == -1]
if len(outers) <= 1:
return [comp]
areas = [(cv2.contourArea(cnts[i]), i) for i in outers]
areas.sort(reverse=True)
biggest_area = areas[0][0]
keep = [i for a, i in areas if a >= 0.15 * biggest_area]
if len(keep) <= 1:
return [comp]
parts = []
for i in keep:
m = np.zeros_like(comp)
cv2.drawContours(m, cnts, i, 1, thickness=cv2.FILLED)
parts.append(m.astype(bool))
return parts
def segment_parts(ink: np.ndarray, cfg: VectorizerConfig,
cb: ProgressCb = None) -> List[Part]:
"""Stage 1: split the ink mask into logical parts.
Default (reliable, offline): connected ink components + conservative
contour-hierarchy split. Optional: MobileSAM refinement (slow, CPU-heavy).
"""
_progress(cb, 0.10, "Stage 1/4 - segmenting parts")
n, labels, stats, _ = cv2.connectedComponentsWithStats(
ink.astype(np.uint8), connectivity=8)
parts: List[Part] = []
pid = 0
for i in range(1, n):
area = stats[i, cv2.CC_STAT_AREA]
if area < cfg.min_part_area:
continue
comp = (labels == i)
for sub in _split_component_hierarchy(comp):
ys, xs = np.nonzero(sub)
if len(xs) == 0:
continue
x0, x1 = xs.min(), xs.max()
y0, y1 = ys.min(), ys.max()
parts.append(Part(id=pid, mask=sub,
bbox=(int(x0), int(y0), int(x1 - x0 + 1), int(y1 - y0 + 1))))
pid += 1
if cfg.use_sam:
sam_parts = _try_sam_parts(ink, cb)
if sam_parts:
parts = sam_parts # SAM wins when it actually runs
_progress(cb, 0.18, f"Stage 1/4 - {len(parts)} parts")
return parts
_MOBILESAM_URL = ("https://github.com/ChaoningZhang/MobileSAM/raw/master/"
"weights/mobile_sam.pt")
def _try_sam_parts(ink: np.ndarray, cb: ProgressCb = None) -> Optional[List[Part]]:
"""Optional MobileSAM automatic mask generation. Returns None on any failure.
NOTE: honest caveat — SAM segments *regions*, not line-art strokes. On
black-outline icons it often returns background blobs, which is why this
is OFF by default and connected components are the reliable path.
"""
try:
from mobile_sam import sam_model_registry, SamAutomaticMaskGenerator
except ImportError:
return None
try:
ckpt = os.path.join(os.path.expanduser("~"), ".cache",
"icon_vectorizer", "mobile_sam.pt")
if not os.path.exists(ckpt):
_progress(cb, 0.12, "Stage 1/4 - downloading MobileSAM weights (one-time)")
os.makedirs(os.path.dirname(ckpt), exist_ok=True)
urllib.request.urlretrieve(_MOBILESAM_URL, ckpt)
_progress(cb, 0.14, "Stage 1/4 - MobileSAM segmenting (slow on CPU)")
sam = sam_model_registry["vit_t"](checkpoint=ckpt)
sam.to("cpu")
gen = SamAutomaticMaskGenerator(sam, points_per_side=16)
rgb = np.stack([(ink * 255).astype(np.uint8)] * 3, axis=-1)
masks = gen.generate(255 - rgb) # SAM expects dark-on-light flipped? keep simple
parts = []
for i, m in enumerate(masks):
seg = m["segmentation"]
if seg.sum() < 200:
continue
ys, xs = np.nonzero(seg)
parts.append(Part(id=i, mask=seg,
bbox=(int(xs.min()), int(ys.min()),
int(xs.max() - xs.min() + 1),
int(ys.max() - ys.min() + 1))))
return parts or None
except Exception:
return None
# ============================================================================
# STAGE 2 - Primitive detection (per part)
# ============================================================================
def _kasa_circle(pts: np.ndarray) -> Tuple[float, float, float]:
"""Algebraic (Kasa) circle fit. pts: Nx2."""
x = pts[:, 0]
y = pts[:, 1]
A = np.c_[2 * x, 2 * y, np.ones(len(x))]
b = x ** 2 + y ** 2
sol, *_ = np.linalg.lstsq(A, b, rcond=None)
cx, cy = float(sol[0]), float(sol[1])
r = math.sqrt(max(0.0, float(sol[2]) + cx ** 2 + cy ** 2))
return cx, cy, r
def ransac_circle(pts: np.ndarray, tol: float = 1.2, iters: int = 80,
min_inlier_ratio: float = 0.70) -> Tuple[bool, Tuple[float, float, float], float]:
"""RANSAC circle refinement. Returns (ok, (cx,cy,r), inlier_ratio)."""
pts = np.asarray(pts, dtype=np.float64)
if len(pts) < 6:
return False, (0.0, 0.0, 0.0), 0.0
best_inliers = np.zeros(len(pts), dtype=bool)
rng = np.random.default_rng(7)
for _ in range(iters):
idx = rng.choice(len(pts), 3, replace=False)
tri = pts[idx]
# skip near-collinear triples
if abs(np.cross(tri[1] - tri[0], tri[2] - tri[0])) < 1e-6:
continue
try:
cx, cy, r = _kasa_circle(tri)
except Exception:
continue
if r < 1.0:
continue
d = np.abs(np.linalg.norm(pts - [cx, cy], axis=1) - r)
inl = d < tol
if inl.sum() > best_inliers.sum():
best_inliers = inl
ratio = float(best_inliers.sum()) / len(pts)
if ratio < min_inlier_ratio or best_inliers.sum() < 6:
return False, (0.0, 0.0, 0.0), ratio
cx, cy, r = _kasa_circle(pts[best_inliers]) # refit on inliers
return True, (cx, cy, r), ratio
def _quad_angles_ok(quad: np.ndarray) -> bool:
"""Check convex quad has ~90 degree corners."""
def ang(a, b, c):
v1 = a - b
v2 = c - b
n = np.linalg.norm(v1) * np.linalg.norm(v2)
if n < 1e-9:
return 0.0
return math.degrees(math.acos(np.clip(np.dot(v1, v2) / n, -1, 1)))
for i in range(4):
a = ang(quad[i], quad[(i + 1) % 4], quad[(i + 2) % 4])
if not (65.0 < a < 115.0):
return False
return True
def _inset_quad(quad: np.ndarray, d: float) -> Optional[np.ndarray]:
"""Offset a convex quad inward by d (parallel edge offset)."""
c = quad.mean(axis=0)
lines = []
for i in range(4):
p1, p2 = quad[i], quad[(i + 1) % 4]
e = p2 - p1
L = np.linalg.norm(e)
if L < 1e-9:
return None
n = np.array([-e[1], e[0]]) / L
if np.dot(n, c - p1) < 0: # flip to inward
n = -n
lines.append((p1 + n * d, e))
out = []
for i in range(4):
p1, e1 = lines[i]
p2, e2 = lines[(i + 1) % 4]
# intersect p1 + t*e1 with p2 + s*e2
M = np.column_stack([e1, -e2])
det = np.linalg.det(M)
if abs(det) < 1e-9:
return None
t = np.linalg.det(np.column_stack([p2 - p1, -e2])) / det
out.append(p1 + t * e1)
return np.array(out)
def detect_primitive(part_mask: np.ndarray, cfg: VectorizerConfig) -> Optional[dict]:
"""Stage 2: circle / rect detection for one part mask (upscaled coords).
Returns dict(kind=..., ...) or None when the part needs centerline tracing.
"""
cnts, _ = cv2.findContours(part_mask.astype(np.uint8), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_NONE)
if not cnts:
return None
cnt = max(cnts, key=cv2.contourArea)
pts = cnt[:, 0, :].astype(np.float64) # x, y
if len(pts) < 8:
return None
area = cv2.contourArea(cnt)
peri = cv2.arcLength(cnt, True)
if peri <= 0:
return None
# --- circle: circularity + RANSAC refinement ---------------------------
circularity = 4 * math.pi * area / (peri * peri)
if cfg.detect_circles and circularity > 0.85:
(cx0, cy0), r0 = cv2.minEnclosingCircle(pts.astype(np.float32))
ok, (cx, cy, r), ratio = ransac_circle(pts, tol=1.2, iters=80,
min_inlier_ratio=0.60)
if not ok: # fall back to min enclosing circle
cx, cy, r, ratio = float(cx0), float(cy0), float(r0), 1.0
# sanity: radius must be plausible vs the part bbox
ys, xs = np.nonzero(part_mask)
bw = max(1.0, float(xs.max() - xs.min()), float(ys.max() - ys.min()))
if 0.35 * bw < r < 0.75 * bw:
return {"kind": "circle", "cx": cx, "cy": cy, "r": r,
"inliers": ratio, "circularity": circularity}
# --- rectangle: convex quad with ~90 deg corners -----------------------
approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)[:, 0, :].astype(np.float64)
if len(approx) == 4 and cv2.isContourConvex(approx.astype(np.float32)):
if _quad_angles_ok(approx):
# filled vs outline: ink/bbox area ratio
x, y, w, h = cv2.boundingRect(cnt)
fill_ratio = float(part_mask.sum()) / max(1.0, float(w * h))
if fill_ratio > 0.70:
# filled rect: output as-is (no centerline inset)
return {"kind": "rect", "pts": approx, "stroke_w": 0.0}
dt = cv2.distanceTransform(part_mask.astype(np.uint8),
cv2.DIST_L2, 3)
stroke_w = 2.0 * float(dt.max()) # upscaled px
inset = _inset_quad(approx, stroke_w / 2.0)
if inset is not None:
return {"kind": "rect", "pts": inset, "stroke_w": stroke_w}
return None
# ============================================================================
# STAGE 3 - Accurate centerline: medial axis + skeleton graph
# ============================================================================
@dataclass
class SkelNode:
id: int
x: float
y: float
kind: str # 'end' | 'junction'
@dataclass
class SkelEdge:
id: int
n1: int # node id (-1 for loop edges)
n2: int
pts: np.ndarray # Nx2 float, x/y in upscaled coords
loop: bool = False
@dataclass
class SkelGraph:
nodes: Dict[int, SkelNode]
edges: Dict[int, SkelEdge]
stroke_w: float # median stroke width, upscaled px
branch_dirs: Dict[Tuple[int, int], Tuple[float, float]] # (node, edge) -> unit dir
def _neighbors8(r: int, c: int, h: int, w: int):
for dr in (-1, 0, 1):
for dc in (-1, 0, 1):
if dr == 0 and dc == 0:
continue
rr, cc = r + dr, c + dc
if 0 <= rr < h and 0 <= cc < w:
yield rr, cc
def _fill_tiny_holes(ink: np.ndarray, max_area: float = 500.0) -> np.ndarray:
"""Fill tiny background holes (raster specks) inside the ink mask.
A 1-2px white speck creates a bogus little skeleton loop; real ring
interiors are orders of magnitude bigger and are left alone.
max_area is in upscaled px^2 (500 ~= 8px diameter at 3x).
"""
try:
from scipy.ndimage import label as _label
except ImportError:
return ink
inv = ~ink
lh, nh = _label(inv, structure=np.ones((3, 3), dtype=int))
if nh == 0:
return ink
border = set(np.unique(np.concatenate([lh[0, :], lh[-1, :],
lh[:, 0], lh[:, -1]])))
out = ink.copy()
for i in range(1, nh + 1):
if i in border:
continue
if int((lh == i).sum()) < max_area:
out[lh == i] = True
return out
def skeleton_graph(ink: np.ndarray) -> Optional[SkelGraph]:
"""Stage 3: skeleton -> pruned node/edge graph.
Uses skimage.morphology.skeletonize (Lee, 1994) — NOT Zhang-Suen — plus
a separate deterministic Euclidean distance transform.
NOTE: the spec asked for medial_axis(return_distance=True), but
skimage 0.26.0's medial_axis is NONDETERMINISTIC (verified: same input
gives 673/696/698/700/701 skeleton pixels across runs, even
single-threaded). A tracer that outputs a different SVG every run is
unacceptable, so we use the deterministic Lee skeletonize + EDT.
Distance-based pruning (dist < 1.2) works exactly the same.
ink: bool array (True = ink), upscaled working space.
"""
if not HAS_SKIMAGE or not HAS_SCIPY:
raise RuntimeError(
"scikit-image and scipy are required (pip install scikit-image scipy)")
ink = _fill_tiny_holes(np.ascontiguousarray(ink))
skel = skeletonize(ink)
dist = distance_transform_edt(ink)
# prune: drop skeleton pixels closer than 1.2px to the boundary
keep = skel & (dist >= 1.2)
if keep.sum() < 4:
keep = skel # very thin stroke: keep everything
if keep.sum() < 4:
return None
h, w = keep.shape
k8 = keep.astype(np.uint8)
nb = cv2.filter2D(k8, -1, np.ones((3, 3), np.uint8)).astype(np.int32) - k8.astype(np.int32)
nb[~keep] = 0
# drop isolated pixels
keep[nb == 0] = False
nb[~keep] = 0
if keep.sum() < 4:
return None
stroke_w = float(2.0 * np.median(dist[keep]))
stroke_w = max(stroke_w, 2.0)
# --- nodes: endpoints (nb==1); junction clusters (nb>2, 8-connected) ---
node_id = np.full((h, w), -1, dtype=np.int32)
nodes: Dict[int, SkelNode] = {}
nid = 0
# junction clusters first
jmask = (nb > 2) & keep
nj, jlabels = cv2.connectedComponents(jmask.astype(np.uint8), connectivity=8)
for j in range(1, nj):
ys, xs = np.nonzero(jlabels == j)
cx = float(xs.mean())
cy = float(ys.mean())
# snap representative pixel to cluster pixel nearest centroid
d2 = (xs - cx) ** 2 + (ys - cy) ** 2
rep = int(np.argmin(d2))
nodes[nid] = SkelNode(id=nid, x=cx, y=cy, kind="junction")
node_id[ys, xs] = nid
nid += 1
# endpoints
ys, xs = np.nonzero((nb == 1) & keep & (node_id < 0))
for y, x in zip(ys, xs):
nodes[nid] = SkelNode(id=nid, x=float(x), y=float(y), kind="end")
node_id[y, x] = nid
nid += 1
node_mask = node_id >= 0
unvis = keep & (~node_mask)
edges: Dict[int, SkelEdge] = {}
eid = 0
def node_xy(i: int) -> Tuple[float, float]:
n = nodes[i]
return (n.x, n.y)
# --- walk edges starting from every node -------------------------------
for i, nd in list(nodes.items()):
# representative pixel(s) of this node
rep_pixels = list(zip(*np.nonzero(node_id == i)))
for (ry, rx) in rep_pixels:
for (qy, qx) in _neighbors8(ry, rx, h, w):
if not unvis[qy, qx]:
continue
# walk
pts = [node_xy(i), (float(qx), float(qy))]
unvis[qy, qx] = False
prev = (qy, qx)
cur = (qy, qx)
end_node = None
while True:
nxt = None
for (ny, nx) in _neighbors8(cur[0], cur[1], h, w):
if (ny, nx) == prev:
continue
if not keep[ny, nx]:
continue
hit = int(node_id[ny, nx])
if hit >= 0 and hit != i:
end_node = hit
break
if hit == i:
continue
if unvis[ny, nx]:
nxt = (ny, nx)
break
if end_node is not None:
pts.append(node_xy(end_node))
break
if nxt is None:
break # dead end: treat as open chain tip
prev, cur = cur, nxt
unvis[cur[0], cur[1]] = False
pts.append((float(cur[1]), float(cur[0])))
edges[eid] = SkelEdge(id=eid, n1=i,
n2=end_node if end_node is not None else -2,
pts=np.array(pts, dtype=np.float64))
eid += 1
# --- leftover pixels: closed loops with no nodes ------------------------
while unvis.any():
ys, xs = np.nonzero(unvis)
sy, sx = int(ys[0]), int(xs[0])
pts = [(float(sx), float(sy))]
unvis[sy, sx] = False
prev = (sy, sx)
cur = (sy, sx)
closed = False
while True:
nxt = None
for (ny, nx) in _neighbors8(cur[0], cur[1], h, w):
if (ny, nx) == prev:
continue
if unvis[ny, nx]:
nxt = (ny, nx)
break
if (ny, nx) == (sy, sx) and len(pts) > 4:
closed = True
break
if closed or nxt is None:
break
prev, cur = cur, nxt
unvis[cur[0], cur[1]] = False
pts.append((float(cur[1]), float(cur[0])))
if len(pts) >= 6:
pts.append(pts[0])
edges[eid] = SkelEdge(id=eid, n1=-1, n2=-1,
pts=np.array(pts, dtype=np.float64), loop=True)
eid += 1
if not edges:
return None
g = SkelGraph(nodes=nodes, edges=edges, stroke_w=stroke_w, branch_dirs={})
# Dissolve pixelation-artifact junctions FIRST (so real tips are never
# mistaken for spurs), then prune true noise branches and collapse
# instability tangles; alternate until stable. Degree-2 junctions are
# always artifacts, never real branches.
for _ in range(4):
_dissolve_degree2(g)
_prune_spurs(g)
_collapse_junction_tangles(g)
_remove_redundant_bridges(g)
_demote_tip_junctions(g)
_snap_endpoints(g)
_dissolve_degree2(g)
_prune_spurs(g)
_collapse_junction_tangles(g)
_remove_redundant_bridges(g)
_fix_degenerate_loops(g)
_measure_branch_dirs(g)
return g
def _edge_length(e: SkelEdge) -> float:
d = np.diff(e.pts, axis=0)
return float(np.linalg.norm(d, axis=1).sum())
def _incident(g: SkelGraph, nid: int) -> List[int]:
return [eid for eid, e in g.edges.items() if e.n1 == nid or e.n2 == nid]
def _prune_spurs(g: SkelGraph):
"""Remove spur branches shorter than 1.2 * global stroke width.
A spur = edge with exactly one endpoint node at a real branch point
(degree >= 3 junction). Uses the GLOBAL stroke width: a branch off a
thick stroke is noise even when the branch itself is thin.
"""
w = g.stroke_w
doomed = []
for eid, e in g.edges.items():
if e.loop or e.n1 < 0 or e.n2 < 0:
continue
k1, k2 = g.nodes[e.n1].kind, g.nodes[e.n2].kind
if (k1 == "end") == (k2 == "end"):
continue # spur = exactly one endpoint
if _edge_length(e) < 1.2 * w:
doomed.append(eid)
for eid in doomed:
e = g.edges.pop(eid)
for nid in (e.n1, e.n2):
if nid >= 0 and g.nodes[nid].kind == "end" and not _incident(g, nid):
g.nodes.pop(nid, None)
def _snap_endpoints(g: SkelGraph):
"""Merge endpoint pairs closer than 0.5 * stroke width (closes tiny gaps)."""
ends = [n for n in g.nodes.values() if n.kind == "end"]
merged = set()
for a in ends:
if a.id in merged:
continue
for b in ends:
if b.id <= a.id or b.id in merged:
continue
d = math.hypot(a.x - b.x, a.y - b.y)
if d < 0.5 * g.stroke_w and d > 1e-9:
# rewire b -> a
mx, my = (a.x + b.x) / 2, (a.y + b.y) / 2
a.x, a.y = mx, my
for e in g.edges.values():
if e.n1 == b.id:
e.n1 = a.id
if e.n2 == b.id:
e.n2 = a.id
g.nodes.pop(b.id, None)
merged.add(b.id)
break
def _dissolve_degree2(g: SkelGraph):
"""Merge the two edges at EVERY degree-2 junction, any angle.
A degree-2 "junction" is never a real branch point — it is a staircase /
pixelation artifact on a continuous curve. Dissolving it keeps a curve as
ONE chain (one cubic) instead of fragments. Sharp corners are NOT lost:
the corner point stays in the merged point list and _split_corners
re-splits there during fitting with corner-preserving tangents.
"""
while True:
dissolved = False
for nid, nd in list(g.nodes.items()):
if nd.kind != "junction":
continue
inc = _incident(g, nid)
if len(inc) != 2:
continue
e1, e2 = g.edges[inc[0]], g.edges[inc[1]]
if e1.loop or e2.loop:
continue
o1 = e1.pts if e1.n1 == nid else e1.pts[::-1] # starts at nid
o2 = e2.pts if e2.n1 == nid else e2.pts[::-1] # starts at nid
merged_pts = np.vstack([o1[::-1], o2[1:]])
new_n1 = e1.n2 if e1.n1 == nid else e1.n1
new_n2 = e2.n2 if e2.n1 == nid else e2.n1
is_loop = (new_n1 == new_n2 and new_n1 >= 0)
g.edges[e1.id] = SkelEdge(
id=e1.id,
n1=-1 if is_loop else new_n1,
n2=-1 if is_loop else new_n2,
pts=merged_pts, loop=is_loop)
g.edges.pop(e2.id)
g.nodes.pop(nid, None)
dissolved = True
break # restart scan: ids changed
if not dissolved:
break
def _fix_degenerate_loops(g: SkelGraph):
"""Edges that start and end at (almost) the same point are loops.
Catches the rare walk artifact: a loop with a single false junction on
it becomes one edge whose tip nearly touches its own tail.
"""
w = g.stroke_w
for e in g.edges.values():
if e.loop or len(e.pts) < 8:
continue
d = float(np.linalg.norm(e.pts[0] - e.pts[-1]))
if d < max(2.0, 0.3 * w) and _edge_length(e) > 2.0 * w:
e.loop = True
e.n1 = e.n2 = -1
# close the point list for loop handling
if d > 1e-9:
e.pts = np.vstack([e.pts, e.pts[0]])
def _collapse_junction_tangles(g: SkelGraph):
"""Merge junctions linked by short junction-junction edges (< 0.8w).
Collapses medial-axis instability tangles — dense little webs of false
junctions where a curve bends — into single junction nodes. Internal tiny
edges are deleted; external edges are rewired to the centroid.
"""
w = g.stroke_w
jids = [nid for nid, n in g.nodes.items() if n.kind == "junction"]
if len(jids) < 2:
return
parent = {j: j for j in jids}
def find(a):
while parent[a] != a:
parent[a] = parent[parent[a]]
a = parent[a]
return a
for e in g.edges.values():
if e.loop or e.n1 < 0 or e.n2 < 0 or e.n1 == e.n2:
continue
if (g.nodes[e.n1].kind == "junction"
and g.nodes[e.n2].kind == "junction"
and _edge_length(e) < 0.8 * w):
ra, rb = find(e.n1), find(e.n2)
if ra != rb:
parent[ra] = rb
groups: Dict[int, List[int]] = {}
for j in jids:
groups.setdefault(find(j), []).append(j)
for members in groups.values():
if len(members) < 2:
continue
new_id = min(members)
nx = sum(g.nodes[m].x for m in members) / len(members)
ny = sum(g.nodes[m].y for m in members) / len(members)
g.nodes[new_id].x, g.nodes[new_id].y = nx, ny
member_set = set(members)
for eid, e in list(g.edges.items()):
if e.loop:
continue
in1 = e.n1 in member_set
in2 = e.n2 in member_set
if in1 and in2:
g.edges.pop(eid) # internal tangle edge
elif in1:
e.n1 = new_id
e.pts[0] = (nx, ny)
elif in2:
e.n2 = new_id
e.pts[-1] = (nx, ny)
for m in members:
if m != new_id:
g.nodes.pop(m, None)
# self-loops created by the merge
for eid, e in list(g.edges.items()):
if not e.loop and e.n1 == new_id and e.n2 == new_id:
if _edge_length(e) < 2.0 * w:
g.edges.pop(eid)
else:
e.loop, e.n1, e.n2 = True, -1, -1
# duplicate edges -> keep the longest
seen: Dict[Tuple[int, int], int] = {}
for eid, e in list(g.edges.items()):
if e.loop or e.n1 < 0 or e.n2 < 0:
continue
key = tuple(sorted((e.n1, e.n2)))
if key in seen:
prev = g.edges[seen[key]]
if _edge_length(e) > _edge_length(prev):
g.edges.pop(seen[key])
seen[key] = eid
else:
g.edges.pop(eid)
else:
seen[key] = eid
def _has_alt_path(g: SkelGraph, src: int, dst: int, exclude: int) -> bool:
"""BFS: is there a path from src to dst not using edge `exclude`?"""
seen = {src}
stack = [src]
while stack:
cur = stack.pop()
if cur == dst:
return True
for eid in _incident(g, cur):
if eid == exclude:
continue
e = g.edges[eid]
nxt = e.n2 if e.n1 == cur else e.n1
if nxt >= 0 and nxt not in seen:
seen.add(nxt)
stack.append(nxt)
return False
def _remove_redundant_bridges(g: SkelGraph):
"""Remove short inter-junction edges that are redundant (a cycle exists).
In dense junction tangles (strokes merging into blobs), the skeleton
grows short cross-edges between nearby junctions. These are artifacts:
removing them doesn't disconnect the graph (an alternative path exists),
but keeping them renders as ugly blobs. A legitimate short stroke (like
the bar of an 'H') has NO alternative path, so it's preserved.
"""
w = g.stroke_w
if w <= 0:
return
for _ in range(10):
removed = False
for eid, e in list(g.edges.items()):
if e.loop or e.n1 < 0 or e.n2 < 0 or e.n1 == e.n2:
continue
n1 = g.nodes.get(e.n1)
n2 = g.nodes.get(e.n2)
if n1 is None or n2 is None:
continue
if n1.kind != "junction" or n2.kind != "junction":
continue
if _edge_length(e) >= 1.0 * w:
continue
if _has_alt_path(g, e.n1, e.n2, exclude=eid):
g.edges.pop(eid, None)
removed = True
break
if not removed:
break
_dissolve_degree2(g)
def _demote_tip_junctions(g: SkelGraph):
"""Junction nodes left with <=1 incident edge are really just tips."""
for nid, nd in list(g.nodes.items()):
if nd.kind == "junction" and len(_incident(g, nid)) <= 1:
nd.kind = "end"
def _measure_branch_dirs(g: SkelGraph):
"""For each junction, measure each branch direction away from the junction
using points 0.8w..2.4w out — used for bump-free tangent constraints."""
w = g.stroke_w
for nid, nd in g.nodes.items():
if nd.kind != "junction":
continue
j = np.array([nd.x, nd.y])
for eid in _incident(g, nid):
e = g.edges[eid]
if e.loop:
continue
p = e.pts if e.n1 == nid else e.pts[::-1]
d = np.linalg.norm(p - j, axis=1)
sel = p[(d >= 0.8 * w) & (d <= 2.4 * w)]
if len(sel) < 2:
sel = p[1:min(len(p), 8)]
if len(sel) < 2:
continue
v = sel.mean(axis=0) - j
n = np.linalg.norm(v)
if n > 1e-9:
g.branch_dirs[(nid, eid)] = (float(v[0] / n), float(v[1] / n))
# ============================================================================
# STAGE 4 - Fitting & semantic labeling
# ============================================================================
def _norm(v: np.ndarray) -> np.ndarray:
n = float(np.linalg.norm(v))
if n < 1e-12:
return v
return v / n
def _dedupe(pts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
if len(pts) < 2:
return pts
keep = [0]
for i in range(1, len(pts)):
if np.linalg.norm(pts[i] - pts[keep[-1]]) > eps:
keep.append(i)
return pts[np.array(keep)]
# ---- Schneider cubic Bezier fit (Graphics Gems, "An Algorithm for
# Automatically Fitting Digitized Curves") ------------------------------
def _chord_param(pts: np.ndarray) -> np.ndarray:
d = np.linalg.norm(np.diff(pts, axis=0), axis=1)
u = np.zeros(len(pts))
u[1:] = np.cumsum(d)
if u[-1] < 1e-12:
return u
return u / u[-1]
def _bez_pt(bez: np.ndarray, t: float) -> np.ndarray:
mt = 1.0 - t
return ((mt ** 3) * bez[0] + 3 * (mt ** 2) * t * bez[1]
+ 3 * mt * (t ** 2) * bez[2] + (t ** 3) * bez[3])
def _gen_bezier(pts: np.ndarray, u: np.ndarray,
tan1: np.ndarray, tan2: np.ndarray) -> np.ndarray:
"""Least-squares solve for alpha1/alpha2. tan1: along travel at start;
tan2: INTO the curve at end (Graphics Gems convention)."""
C = np.zeros((2, 2))
X = np.zeros(2)
p0, p3 = pts[0], pts[-1]
for i in range(len(pts)):
t = u[i]
mt = 1.0 - t
b1 = 3 * t * mt * mt
b2 = 3 * t * t * mt
a1 = tan1 * b1
a2 = tan2 * b2
C[0, 0] += float(a1 @ a1)
C[0, 1] += float(a1 @ a2)
C[1, 1] += float(a2 @ a2)
tmp = pts[i] - (p0 * (mt ** 3) + p0 * b1 + p3 * b2 + p3 * (t ** 3))
X[0] += float(a1 @ tmp)
X[1] += float(a2 @ tmp)
C[1, 0] = C[0, 1]
det = C[0, 0] * C[1, 1] - C[0, 1] * C[1, 0]
seg_len = float(np.linalg.norm(p3 - p0))
if abs(det) < 1e-12:
al1 = al2 = seg_len / 3.0
else:
al1 = (X[0] * C[1, 1] - X[1] * C[0, 1]) / det
al2 = (C[0, 0] * X[1] - C[1, 0] * X[0]) / det
eps = 1e-6 * max(seg_len, 1.0)
if al1 < eps or al2 < eps: # Wu/Barsky heuristic fallback
al1 = al2 = seg_len / 3.0
return np.array([p0, p0 + tan1 * al1, p3 + tan2 * al2, p3])
def _max_err(pts: np.ndarray, bez: np.ndarray,
u: np.ndarray) -> Tuple[float, int]:
maxd = 0.0
idx = len(pts) // 2
for i in range(len(pts)):
d = float(np.linalg.norm(_bez_pt(bez, u[i]) - pts[i]))
if d > maxd:
maxd, idx = d, i
return maxd, idx
def _newton(bez: np.ndarray, p: np.ndarray, t: float) -> float:
q = _bez_pt(bez, t)
mt = 1.0 - t
d1 = (3 * mt * mt * (bez[1] - bez[0]) + 6 * mt * t * (bez[2] - bez[1])
+ 3 * t * t * (bez[3] - bez[2]))
d2 = 6 * mt * (bez[2] - 2 * bez[1] + bez[0]) + 6 * t * (bez[3] - 2 * bez[2] + bez[1])
num = float((q - p) @ d1)
den = float(d1 @ d1 + (q - p) @ d2)
if abs(den) < 1e-12:
return t
return t - num / den
def _reparam(bez: np.ndarray, pts: np.ndarray, u: np.ndarray) -> np.ndarray:
out = np.empty_like(u)
for i in range(len(pts)):
out[i] = _newton(bez, pts[i], u[i])
return np.clip(out, 0.0, 1.0)
def _fit_cubic(pts: np.ndarray, tan1: np.ndarray, tan2: np.ndarray,
tol: float, depth: int = 0) -> List[np.ndarray]:
u = _chord_param(pts)
bez = _gen_bezier(pts, u, tan1, tan2)
err, idx = _max_err(pts, bez, u)
if err >= tol:
# Newton-Raphson reparameterization BEFORE splitting: chord-length
# is only an initial guess; refining it often avoids a split.
# (Graphics Gems tries this when err < tol^2; we try it always —
# it can only improve the fit and costs 4 iterations.)
for _ in range(4):
u = _reparam(bez, pts, u)
bez = _gen_bezier(pts, u, tan1, tan2)
err, idx = _max_err(pts, bez, u)
if err < tol:
break
if err < tol:
return [bez]
if depth >= 12 or len(pts) <= 4:
return [bez]
idx = max(2, min(len(pts) - 3, idx))
tc = _norm(pts[idx - 1] - pts[idx + 1]) # into-curve tangent at split
return (_fit_cubic(pts[:idx + 1], tan1, tc, tol, depth + 1)
+ _fit_cubic(pts[idx:], -tc, tan2, tol, depth + 1))
def schneider_fit(pts: np.ndarray, tan1: np.ndarray, tan2: np.ndarray,
tol: float) -> List[np.ndarray]:
"""Public wrapper: fit digitized points with cubic Beziers within tol."""
pts = _dedupe(np.asarray(pts, dtype=np.float64))
if len(pts) < 2:
return []
if len(pts) == 2:
p0, p3 = pts
d = p3 - p0
return [np.array([p0, p0 + d / 3, p0 + 2 * d / 3, p3])]
return _fit_cubic(pts, _norm(tan1), _norm(tan2), tol)
def try_single_cubic(pts: np.ndarray, tan1: np.ndarray, tan2: np.ndarray,
tol: float) -> Optional[List[np.ndarray]]:
"""Try to fit ALL points with ONE cubic within tol.
Uses Newton-Raphson reparameterization to get the best single cubic.
Returns [bezier] on success, None if a single cubic can't hold tol.
"""
pts = _dedupe(np.asarray(pts, dtype=np.float64))
if len(pts) < 3:
return None
tan1, tan2 = _norm(tan1), _norm(tan2)
u = _chord_param(pts)
bez = _gen_bezier(pts, u, tan1, tan2)
for _ in range(6):
u = _reparam(bez, pts, u)
bez = _gen_bezier(pts, u, tan1, tan2)
err, _ = _max_err(pts, bez, u)
return [bez] if err < tol else None
# ---- deviation measurement ---------------------------------------------------
def _polyline_dev(pts: np.ndarray, poly: np.ndarray) -> Tuple[float, float]:
"""max/mean distance from pts to a polyline (Nx2). Vectorized."""
if len(poly) < 2 or len(pts) == 0:
return 0.0, 0.0
a = poly[:-1][None, :, :] # 1 x M x 2
b = poly[1:][None, :, :]
p = pts[:, None, :] # N x 1 x 2
ab = b - a
denom = np.sum(ab ** 2, axis=2, keepdims=True) + 1e-12
t = np.clip(np.sum((p - a) * ab, axis=2, keepdims=True) / denom, 0, 1)
proj = a + t * ab
d = np.linalg.norm(p - proj, axis=2)
dmin = d.min(axis=1)
return float(dmin.max()), float(dmin.mean())
def _sample_beziers(beziers: List[np.ndarray], per: int = 24) -> np.ndarray:
out = []
for bez in beziers:
ts = np.linspace(0, 1, per)
for t in ts:
out.append(_bez_pt(bez, t))
return np.array(out)
# ---- chain subdivision: corners / straight / arc ------------------------------
def _split_corners(pts: np.ndarray, tan1: np.ndarray, tan2: np.ndarray,
stroke_w: float) -> List[Tuple[int, int, np.ndarray, np.ndarray]]:
"""Split a chain at sharp corners (>55 deg turn). Returns
[(a, b, tan1, tan2), ...] index ranges with corner-preserving tangents,
so callers can slice matching raw (unsmoothed) points for dev stats."""
k = max(3, int(round(stroke_w / 2)))
corners = []
for i in range(k, len(pts) - k):
v1 = pts[i] - pts[i - k]
v2 = pts[i + k] - pts[i]
n1, n2 = np.linalg.norm(v1), np.linalg.norm(v2)
if n1 < 1e-9 or n2 < 1e-9:
continue
turn = math.degrees(math.acos(np.clip(float(v1 @ v2) / (n1 * n2), -1, 1)))
if turn > 55:
corners.append(i)
if not corners:
return [(0, len(pts) - 1, tan1, tan2)]
# keep corners separated by at least k
kept = [corners[0]]
for c in corners[1:]:
if c - kept[-1] >= k:
kept.append(c)
out = []
bounds = [0] + kept + [len(pts) - 1]
for s in range(len(bounds) - 1):
a, b = bounds[s], bounds[s + 1]
t1 = tan1 if s == 0 else _norm(pts[a + 1] - pts[a - 1])
t2 = tan2 if s == len(bounds) - 2 else _norm(pts[b - 1] - pts[b + 1])
out.append((a, b, t1, t2))
return out
def _is_straight(pts: np.ndarray, tol: float) -> bool:
if len(pts) < 3:
return True
chord = pts[-1] - pts[0]
L = float(np.linalg.norm(chord))
if L < 1e-9:
return False
n = np.array([-chord[1], chord[0]]) / L
return bool(np.abs((pts - pts[0]) @ n).max() < tol * 0.75)
def _try_arc(pts: np.ndarray, tol: float) -> Optional[dict]:
"""Kasa circle fit on an open chain -> SVG arc params."""
if len(pts) < 6:
return None
try:
cx, cy, r = _kasa_circle(pts)
except Exception:
return None
if r < 2.0 or not np.isfinite(r):
return None
dev = np.abs(np.linalg.norm(pts - [cx, cy], axis=1) - r)
if dev.max() > tol:
return None
angs = np.arctan2(pts[:, 1] - cy, pts[:, 0] - cx) # y-down: + = clockwise
diffs = np.diff(angs)
diffs = (diffs + math.pi) % (2 * math.pi) - math.pi
span = float(np.degrees(diffs.sum()))
if abs(span) < 5:
return None
return {"cx": cx, "cy": cy, "r": r,
"large_arc": 1 if abs(span) > 180 else 0,
"sweep": 1 if span > 0 else 0,
"x1": float(pts[0, 0]), "y1": float(pts[0, 1]),
"x2": float(pts[-1, 0]), "y2": float(pts[-1, 1])}
def _break_loop(pts: np.ndarray) -> np.ndarray:
"""Open a closed loop at its sharpest-curvature point."""
n = len(pts)
k = max(2, n // 36)
best, bi = -1.0, 0
for i in range(n):
v1 = pts[i] - pts[(i - k) % n]
v2 = pts[(i + k) % n] - pts[i]
n1, n2 = np.linalg.norm(v1), np.linalg.norm(v2)
if n1 < 1e-9 or n2 < 1e-9:
continue
turn = math.degrees(math.acos(np.clip(float(v1 @ v2) / (n1 * n2), -1, 1)))
if turn > best:
best, bi = turn, i
return np.roll(pts, -bi, axis=0)
def _reduce_anchors(beziers: List[np.ndarray], pts: np.ndarray,
tol: float) -> List[np.ndarray]:
"""Greedy merge of consecutive cubics that still fit within tol*1.5."""
if len(beziers) < 2:
return beziers
out = [beziers[0]]
# sample points per bezier for refit attempts
for bez in beziers[1:]:
cand = _sample_beziers([out[-1], bez], per=16)
t1 = _norm(cand[1] - cand[0])
t2 = _norm(cand[-2] - cand[-1])
merged = _fit_cubic(cand, t1, t2, tol * 1.5, depth=0)
if len(merged) == 1:
err, _ = _max_err(cand, merged[0], _chord_param(cand))
if err < tol * 1.5:
out[-1] = merged[0]
continue
out.append(bez)
return out
# ---- per-edge fitting ----------------------------------------------------------
def _estimate_tangent(pts: np.ndarray, at_start: bool) -> np.ndarray:
k = min(4, len(pts) - 1)
if at_start:
return _norm(pts[k] - pts[0])
return _norm(pts[-k - 1] - pts[-1]) # into-curve at end
def _fit_edge(e: SkelEdge, g: SkelGraph, cfg: VectorizerConfig,
scale: int, base: str, part_idx: int) -> List[dict]:
fp = cfg.fit_params()
tol_w = fp.tol * scale # working (upscaled) pixels
pts = _dedupe(e.pts.astype(np.float64))
if len(pts) < 2:
return []
closed = False
# --- closed loops: circle -> ellipse -> break open -----------------------
if e.loop:
ok, (cx, cy, r), ratio = ransac_circle(pts, tol=1.2, iters=80,
min_inlier_ratio=0.90)
if ok:
return [_circle_el(cx / scale, cy / scale, r / scale, base,
part_idx, g.stroke_w / scale)]
if len(pts) >= 12:
try:
(ecx, ecy), (rw, rh), ang = cv2.fitEllipse(pts.astype(np.float32))
rx, ry = rw / 2, rh / 2
if min(rx, ry) > 2:
th = np.linspace(0, 2 * math.pi, 72)
ex = ecx + rx * np.cos(th) * math.cos(math.radians(ang)) \
- ry * np.sin(th) * math.sin(math.radians(ang))
ey = ecy + rx * np.cos(th) * math.sin(math.radians(ang)) \
+ ry * np.sin(th) * math.cos(math.radians(ang))
samp = np.c_[ex, ey]
md, _ = _polyline_dev(pts, samp)
if md < fp.tol * 2.5:
return [_ellipse_el(ecx / scale, ecy / scale,
rx / scale, ry / scale, ang,
base, part_idx, g.stroke_w / scale)]
except Exception:
pass
pts = _break_loop(pts)
closed = True
raw = pts.copy()
# --- light pre-smoothing (Low = 1 pass) ----------------------------------
if HAS_SCIPY and fp.smooth_iters > 0 and len(pts) > 4:
for _ in range(fp.smooth_iters):
pts[:, 0] = gaussian_filter1d(pts[:, 0], SMOOTH_SIGMA,
mode="nearest")
pts[:, 1] = gaussian_filter1d(pts[:, 1], SMOOTH_SIGMA,
mode="nearest")
# --- endpoint tangents: junction branch dirs = no bumps -------------------
if not e.loop and e.n1 >= 0 and g.nodes[e.n1].kind == "junction" \
and cfg.detect_junctions and (e.n1, e.id) in g.branch_dirs:
tan1 = np.array(g.branch_dirs[(e.n1, e.id)])
else:
tan1 = _estimate_tangent(pts, True)
if not e.loop and e.n2 >= 0 and g.nodes[e.n2].kind == "junction" \
and cfg.detect_junctions and (e.n2, e.id) in g.branch_dirs:
tan2 = -np.array(g.branch_dirs[(e.n2, e.id)]) # into-curve
else:
tan2 = _estimate_tangent(pts, False)
chains = (_split_corners(pts, tan1, tan2, g.stroke_w)
if cfg.detect_corners else [(0, len(pts) - 1, tan1, tan2)])
els: List[dict] = []
sw = g.stroke_w / scale
for a, b, t1, t2 in chains:
c = _dedupe(pts[a:b + 1])
rc = _dedupe(raw[a:b + 1]) # matching unsmoothed points for dev stats
if len(c) < 2:
continue
if _is_straight(c, tol_w):
els.append(_line_el(c, rc, scale, base, part_idx, sw, closed and len(chains) == 1))
continue
arc = _try_arc(c, tol_w) if cfg.detect_circles else None
if arc is not None:
els.append(_arc_el(arc, rc, scale, base, part_idx, sw))
continue
# Prefer ONE cubic for smooth chains (not 20 fragments): try a
# single cubic first with a relaxed tolerance; fall back to the
# accurate recursive fit when the shape genuinely needs it.
single = try_single_cubic(c, t1, t2, tol_w * SINGLE_CUBIC_TOL_FACTOR)
if single is not None:
beziers = single
else:
beziers = schneider_fit(c, t1, t2, tol_w)
if fp.reduce:
beziers = _reduce_anchors(beziers, c, tol_w)
els.append(_path_el(beziers, rc, scale, base, part_idx, sw,
closed and len(chains) == 1))
return els
# ---- element constructors (original-image coords) ------------------------------
def _dev_stats(raw_up: np.ndarray, sampled_up: np.ndarray, scale: int):
md, mn = _polyline_dev(raw_up, sampled_up)
return md / scale, mn / scale
def _circle_el(cx, cy, r, base, part_idx, sw):
return {"kind": "circle", "cx": cx, "cy": cy, "r": r, "base": base,
"part": part_idx, "anchors": 4, "max_dev": 0.0, "mean_dev": 0.0,
"stroke_w": sw, "name": ""}
def _ellipse_el(cx, cy, rx, ry, ang, base, part_idx, sw):
return {"kind": "ellipse", "cx": cx, "cy": cy, "rx": rx, "ry": ry,
"angle": ang, "base": base, "part": part_idx, "anchors": 4,
"max_dev": 0.0, "mean_dev": 0.0, "stroke_w": sw, "name": ""}
def _line_el(c_up, raw_up, scale, base, part_idx, sw, closed):
p0, p1 = c_up[0] / scale, c_up[-1] / scale
md, mn = _dev_stats(raw_up, c_up, scale)
return {"kind": "line", "p0": p0, "p1": p1, "base": base, "part": part_idx,
"anchors": 2, "max_dev": md, "mean_dev": mn, "stroke_w": sw,
"name": "", "closed": closed}
def _arc_el(arc, raw_up, scale, base, part_idx, sw):
cx, cy, r = arc["cx"] / scale, arc["cy"] / scale, arc["r"] / scale
th = np.linspace(0, 1, 33)
# sample along the arc for dev stats
a0 = math.atan2(arc["y1"] - arc["cy"], arc["x1"] - arc["cx"])
span = math.radians(360)
# reconstruct signed span from flags (approx for stats only)
a1 = math.atan2(arc["y2"] - arc["cy"], arc["x2"] - arc["cx"])
d = (a1 - a0)
while d <= -math.pi:
d += 2 * math.pi
while d > math.pi:
d -= 2 * math.pi
if arc["sweep"] == 0 and d > 0:
d -= 2 * math.pi
if arc["sweep"] == 1 and d < 0:
d += 2 * math.pi
angs = a0 + d * th
samp = np.c_[arc["cx"] + arc["r"] * np.cos(angs),
arc["cy"] + arc["r"] * np.sin(angs)]
md, mn = _dev_stats(raw_up, samp, scale)
return {"kind": "arc", "cx": cx, "cy": cy, "r": r,
"large_arc": arc["large_arc"], "sweep": arc["sweep"],
"x1": arc["x1"] / scale, "y1": arc["y1"] / scale,
"x2": arc["x2"] / scale, "y2": arc["y2"] / scale,
"base": base, "part": part_idx, "anchors": 3,
"max_dev": md, "mean_dev": mn, "stroke_w": sw, "name": ""}
def _path_el(beziers_up, raw_up, scale, base, part_idx, sw, closed):
beziers = [b / scale for b in beziers_up]
samp = _sample_beziers(beziers_up)
md, mn = _dev_stats(raw_up, samp, scale)
return {"kind": "path", "beziers": beziers, "base": base, "part": part_idx,
"anchors": 3 * len(beziers) + 1, "max_dev": md, "mean_dev": mn,
"stroke_w": sw, "name": "", "closed": closed}
def _element_from_primitive(prim: dict, scale: int, base: str, part_idx: int) -> dict:
sw = 0.0
if prim["kind"] == "circle":
sw = 2.0 * prim["r"] / scale * 0.12 # placeholder, refined below
el = _circle_el(prim["cx"] / scale, prim["cy"] / scale,
prim["r"] / scale, base, part_idx, sw)
el["prim_inliers"] = prim.get("inliers", 1.0)
return el
# rect: inset quad in original coords
q = prim["pts"] / scale
el = {"kind": "rect", "quad": q, "base": base, "part": part_idx,
"anchors": 4, "max_dev": 0.0, "mean_dev": 0.0,
"stroke_w": prim.get("stroke_w", 0.0) / scale, "name": ""}
return el
# ============================================================================
# CLIP semantic labeling (optional, offline after one-time weight download)
# ============================================================================
_clip_cache: dict = {}
def clip_label_parts(part_crops: List[Image.Image],
prompts: List[str]) -> Optional[List[str]]:
"""Label each part crop with CLIP ViT-B/32. Returns None on any failure
(missing torch/models -> caller falls back to generic names)."""
if not part_crops or not prompts:
return None
try:
import torch
try:
import open_clip
if "open_clip" not in _clip_cache:
model, _, preprocess = open_clip.create_model_and_transforms(
"ViT-B-32", pretrained="openai")
tok = open_clip.get_tokenizer("ViT-B-32")
_clip_cache["open_clip"] = (model, preprocess, tok, "oc")
model, preprocess, tok, _ = _clip_cache["open_clip"]
device = "cpu"
model.to(device).eval()
text = tok([f"an icon of {p}" for p in prompts]).to(device)
with torch.no_grad():
tf = model.encode_text(text)
tf = tf / tf.norm(dim=-1, keepdim=True)
labels = []
for crop in part_crops:
img = preprocess(crop).unsqueeze(0).to(device)
imf = model.encode_image(img)
imf = imf / imf.norm(dim=-1, keepdim=True)
sim = (100.0 * imf @ tf.T).softmax(dim=-1)
labels.append(prompts[int(sim.argmax())])
return labels
except ImportError:
pass
import clip as oclip # OpenAI's clip package
if "oa_clip" not in _clip_cache:
model, preprocess = oclip.load("ViT-B/32", device="cpu")
_clip_cache["oa_clip"] = (model, preprocess)
model, preprocess = _clip_cache["oa_clip"]
model.eval()
text = oclip.tokenize([f"an icon of {p}" for p in prompts])
with torch.no_grad():
tf = model.encode_text(text)
tf = tf / tf.norm(dim=-1, keepdim=True)
labels = []
for crop in part_crops:
img = preprocess(crop).unsqueeze(0)
imf = model.encode_image(img)
imf = imf / imf.norm(dim=-1, keepdim=True)
sim = (100.0 * imf @ tf.T).softmax(dim=-1)
labels.append(prompts[int(sim.argmax())])
return labels
except Exception:
return None
def _label_parts(gray: np.ndarray, parts: List[Part], scale: int,
cfg: VectorizerConfig) -> List[str]:
"""Crop each part from the original and label with CLIP (fallback: '')."""
if not cfg.use_clip or not parts:
return [""] * len(parts)
H, W = gray.shape[:2]
rgb = Image.fromarray(np.stack([gray] * 3, axis=-1))
crops = []
for p in parts:
x, y, w, h = p.bbox
x0, y0 = max(0, x // scale - 4), max(0, y // scale - 4)
x1, y1 = min(W, (x + w) // scale + 4), min(H, (y + h) // scale + 4)
crop = rgb.crop((x0, y0, max(x1, x0 + 8), max(y1, y0 + 8)))
side = max(crop.size)
sq = Image.new("RGB", (side, side), (255, 255, 255))
sq.paste(crop, ((side - crop.size[0]) // 2, (side - crop.size[1]) // 2))
crops.append(sq.resize((224, 224), Image.BICUBIC))
labels = clip_label_parts(crops, cfg.clip_prompts)
if not labels:
return [""] * len(parts)
return [l.strip().lower().replace(" ", "-") for l in labels]
def _assign_names(elements: List[dict]):
"""data-name assignment: circle -> {base}-rim, else {base}; dedupe."""
used: Dict[str, int] = {}
for el in elements:
base = el.get("base", "")
if base:
key = f"{base}-rim" if el["kind"] == "circle" else base
else:
key = {"circle": "circle", "ellipse": "ellipse",
"rect": "rect"}.get(el["kind"], "path")
n = used.get(key, 0) + 1
used[key] = n
el["name"] = key if n == 1 else f"{key}-{n}"
# ============================================================================
# SVG writer
# ============================================================================
def _f(v: float) -> str:
s = f"{v:.1f}"
return s
def _path_d(beziers: List[np.ndarray], closed: bool) -> str:
p0 = beziers[0][0]
d = f"M{_f(p0[0])} {_f(p0[1])}"
for b in beziers:
d += (f"C{_f(b[1][0])} {_f(b[1][1])} {_f(b[2][0])} {_f(b[2][1])} "
f"{_f(b[3][0])} {_f(b[3][1])}")
if closed:
d += "Z"
return d
def _element_svg(el: dict, indent: str) -> str:
n = el["name"]
k = el["kind"]
if k == "circle":
return (f'{indent}<circle data-name="{n}" cx="{_f(el["cx"])}" '
f'cy="{_f(el["cy"])}" r="{_f(el["r"])}"/>')
if k == "ellipse":
t = (f' transform="rotate({_f(el["angle"])} {_f(el["cx"])} {_f(el["cy"])})"'
if abs(el["angle"]) > 0.5 else "")
return (f'{indent}<ellipse data-name="{n}" cx="{_f(el["cx"])}" '
f'cy="{_f(el["cy"])}" rx="{_f(el["rx"])}" ry="{_f(el["ry"])}"{t}/>')
if k == "rect":
q = el["quad"]
xs, ys = q[:, 0], q[:, 1]
# axis-aligned (order-agnostic): every corner's x is ~x0/x1 and y ~y0/y1
x0, x1 = float(xs.min()), float(xs.max())
y0, y1 = float(ys.min()), float(ys.max())
aa = True
for x, y in q:
if min(abs(x - x0), abs(x - x1)) > 1.0:
aa = False
break
if min(abs(y - y0), abs(y - y1)) > 1.0:
aa = False
break
if aa:
return (f'{indent}<rect data-name="{n}" x="{_f(x0)}" y="{_f(y0)}" '
f'width="{_f(x1 - x0)}" height="{_f(y1 - y0)}"/>')
d = f"M{_f(xs[0])} {_f(ys[0])}"
for i in range(1, 4):
d += f"L{_f(xs[i])} {_f(ys[i])}"
return f'{indent}<path data-name="{n}" d="{d}Z"/>'
if k == "line":
p0, p1 = el["p0"], el["p1"]
d = f"M{_f(p0[0])} {_f(p0[1])}L{_f(p1[0])} {_f(p1[1])}"
if el.get("closed"):
d += "Z"
return f'{indent}<path data-name="{n}" d="{d}"/>'
if k == "arc":
d = (f'M{_f(el["x1"])} {_f(el["y1"])}A{_f(el["r"])} {_f(el["r"])} 0 '
f'{el["large_arc"]},{el["sweep"]} {_f(el["x2"])} {_f(el["y2"])}')
return f'{indent}<path data-name="{n}" d="{d}"/>'
# path
d = _path_d(el["beziers"], el.get("closed", False))
return f'{indent}<path data-name="{n}" d="{d}"/>'
def build_svg(elements: List[dict], W: int, H: int, stroke_w: float,
title: str = "icon") -> str:
lines = ['<?xml version="1.0" encoding="UTF-8"?>',
f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {W} {H}" '
f'width="{W}" height="{H}">',
f' <g fill="none" stroke="black" stroke-width="{_f(stroke_w)}">']
# group elements sharing a CLIP base label (like the eyepiece example)
groups: Dict[str, List[dict]] = {}
order: List[str] = []
for el in elements:
base = el.get("base", "")
key = base if base else f"\x00{el['name']}"
if key not in groups:
groups[key] = []
order.append(key)
groups[key].append(el)
for key in order:
els = groups[key]
if key.startswith("\x00") or len(els) == 1:
for el in els:
lines.append(_element_svg(el, " "))
else:
lines.append(f' <g data-name="{key}">')
for el in els:
lines.append(_element_svg(el, " "))
lines.append(" </g>")
lines += [" </g>", "</svg>"]
return "\n".join(lines) + "\n"
# ============================================================================
# Main entry points
# ============================================================================
def convert_image_to_svg(image_path: str, cfg: VectorizerConfig,
cb: ProgressCb = None) -> dict:
"""Full 4-stage pipeline. Returns dict(svg, elements, stats, W, H, ...)."""
t0 = time.time()
gray = load_grayscale(image_path)
H, W = gray.shape[:2]
_progress(cb, 0.02, "Stage 0/4 - preprocessing")
ink, scale = preprocess(gray, cfg)
# Dual threshold: segment parts at a LOWER threshold so antialiased gray
# halos don't merge separate strokes into one blob (topology), while
# stroke width is measured from the main threshold mask (geometry).
seg_th = max(50, int(cfg.threshold) - 30)
if seg_th < int(cfg.threshold):
ink_seg, _ = preprocess(gray, replace(cfg, threshold=seg_th))
else:
ink_seg = ink
parts = segment_parts(ink_seg, cfg, cb)
# Distance transform of the main-threshold mask: used to measure the
# true stroke width under each (thin) segmentation part.
try:
from scipy.ndimage import distance_transform_edt as _edt
_dt_hi = _edt(ink)
except Exception:
_dt_hi = None
_progress(cb, 0.19, "Stage 4/4 - CLIP labeling")
labels = _label_parts(gray, parts, scale, cfg)
elements: List[dict] = []
sw_list: List[float] = []
n_parts = max(1, len(parts))
for idx, part in enumerate(parts):
_progress(cb, 0.20 + 0.70 * idx / n_parts,
f"Stages 2-3/4 - part {idx + 1}/{len(parts)}")
base = labels[idx] if idx < len(labels) else ""
prim = detect_primitive(part.mask, cfg)
if prim is not None:
el = _element_from_primitive(prim, scale, base, idx)
sw_list.append(el["stroke_w"] if el["stroke_w"] > 0
else 3.0 / scale)
# refine circle stroke width from local distance transform
elements.append(el)
continue
g = skeleton_graph(part.mask)
if g is None:
continue
sw_list.append(g.stroke_w / scale)
for e in g.edges.values():
try:
elements.extend(_fit_edge(e, g, cfg, scale, base, idx))
except Exception:
continue
_progress(cb, 0.94, "Writing SVG")
_assign_names(elements)
# Robust global stroke width: for each (thin) segmentation part, measure
# the full stroke width from the main-threshold distance transform, then
# take the 75th percentile (tiny detail parts don't drag it down).
# Robust to merged blobs and thin-mask underestimation.
_widths = []
if _dt_hi is not None:
for part in parts:
try:
d = _dt_hi[part.mask]
if d.size > 100:
_widths.append(2.0 * float(np.percentile(d, 90)) / scale)
except Exception:
pass
if _widths:
stroke_w = float(np.percentile(_widths, 75))
else:
stroke_w = float(np.median(sw_list)) if sw_list else 2.0
# fix primitive-circle stroke widths with the global median
for el in elements:
if el["kind"] == "circle" and el.get("prim_inliers") is not None:
el["stroke_w"] = stroke_w
svg = build_svg(elements, W, H, stroke_w)
max_dev = max([e["max_dev"] for e in elements], default=0.0)
mean_dev = (sum(e["mean_dev"] for e in elements) / len(elements)
if elements else 0.0)
stats = {
"paths": len(elements),
"anchors": sum(e["anchors"] for e in elements),
"stroke_w": round(stroke_w, 1),
"max_dev": round(max_dev, 2),
"mean_dev": round(mean_dev, 2),
"parts": len(parts),
"elapsed": round(time.time() - t0, 1),
}
_progress(cb, 1.0, f"Done in {stats['elapsed']}s")
return {"svg": svg, "elements": elements, "stats": stats,
"W": W, "H": H, "scale": scale}
def convert_file(input_path: str, output_path: Optional[str] = None,
cfg: Optional[VectorizerConfig] = None,
cb: ProgressCb = None) -> Tuple[str, dict]:
"""Convert one image; default output: <input_dir>/output_svgs/<name>.svg."""
cfg = cfg or VectorizerConfig()
if output_path is None:
d = os.path.join(os.path.dirname(os.path.abspath(input_path)),
"output_svgs")
os.makedirs(d, exist_ok=True)
stem = os.path.splitext(os.path.basename(input_path))[0]
output_path = os.path.join(d, stem + ".svg")
res = convert_image_to_svg(input_path, cfg, cb)
with open(output_path, "w", encoding="utf-8") as f:
f.write(res["svg"])
return output_path, res
if __name__ == "__main__":
import argparse
ap = argparse.ArgumentParser(description="ICON VECTORIZER - ACCURATE CENTERLINE")
ap.add_argument("input", help="input JPG/PNG icon")
ap.add_argument("-o", "--output", default=None)
ap.add_argument("--threshold", type=int, default=110)
ap.add_argument("--upscale", type=int, default=3, choices=[1, 2, 3, 4])
ap.add_argument("--smoothing", default="Low", choices=["Low", "Medium", "High"])
ap.add_argument("--tolerance", type=float, default=0.6)
ap.add_argument("--no-corners", action="store_true")
ap.add_argument("--no-junctions", action="store_true")
ap.add_argument("--no-circles", action="store_true")
ap.add_argument("--reduce-anchors", action="store_true")
ap.add_argument("--no-clip", action="store_true")
ap.add_argument("--sam", action="store_true")
a = ap.parse_args()
cfg = VectorizerConfig(
threshold=a.threshold, upscale=a.upscale, smoothing=a.smoothing,
tolerance=a.tolerance, detect_corners=not a.no_corners,
detect_junctions=not a.no_junctions, detect_circles=not a.no_circles,
reduce_anchors=a.reduce_anchors, use_clip=not a.no_clip,
use_sam=a.sam)
out, res = convert_file(a.input, a.output, cfg,
cb=lambda f, s: print(f"\r{f*100:5.1f}% {s}", end=""))
print(f"\n-> {out}")
print(res["stats"])
for el in res["elements"]:
print(f' {el["kind"]:7s} {el["name"]:24s} anchors={el["anchors"]} '
f'maxdev={el["max_dev"]:.2f}')
"""
ICON VECTORIZER - ACCURATE CENTERLINE
====================================
vectorizer_core.py — slow-but-accurate 4-stage centerline vectorization pipeline.
Converts black-outline icons (JPG/PNG) into clean centerline SVG with semantic
data-name labels (e.g. eyepiece-rim, optical-tube, curved-arm).
Pipeline
--------
Stage 0 - Preprocessing : 3x/4x INTER_CUBIC upscale, median+gaussian denoise,
user threshold, morphological opening.
Stage 1 - Segmentation : connected ink components (+ contour-hierarchy split),
optional MobileSAM refinement.
Stage 2 - Primitives : circularity test -> perfect <circle> (minEnclosingCircle
+ RANSAC refine); convex quad -> inset <rect>.
Stage 3 - Centerline : skimage skeletonize (Lee, NOT Zhang-Suen) + EDT distance,
skeleton graph (endpoints/junctions/loops), spur prune,
endpoint snap, junction branch-direction measurement.
Stage 4 - Fitting : RANSAC circle/ellipse on closed loops, Kasa arc fit,
Schneider cubic Bezier fit (Graphics Gems, tol=0.6px),
optional CLIP semantic labeling, clean SVG output.
Everything geometric is deterministic and offline. Only the OPTIONAL semantic
labeling (CLIP) and segmentation refinement (MobileSAM) need torch + model
downloads; the tracer works fully without them.
Author: built for Twins Design Studio / allfa — accuracy >> speed.
"""
from __future__ import annotations
import math
import os
import time
import urllib.request
from dataclasses import dataclass, field, replace
from typing import Callable, Dict, List, Optional, Tuple
import numpy as np
import cv2
from PIL import Image
try:
from skimage.morphology import skeletonize
HAS_SKIMAGE = True
except ImportError: # pragma: no cover
HAS_SKIMAGE = False
try:
from scipy.ndimage import gaussian_filter1d, distance_transform_edt
HAS_SCIPY = True
except ImportError: # pragma: no cover
HAS_SCIPY = False
# ----------------------------------------------------------------------------
# Configuration
# ----------------------------------------------------------------------------
SMOOTH_PRESETS = {"Low": 1, "Medium": 4, "High": 10}
class FitParams:
"""Fitting parameters (accurate mode).
tol is in ORIGINAL image pixels (user-facing unit); the pipeline works
at `scale`x, so the working tolerance is tol * scale.
"""
tol: float = 0.6 # max deviation in px (original image size)
smooth_iters: int = 1 # light pre-smoothing passes (Low)
reduce: bool = False # merge near-collinear cubics (OFF for accuracy)
# Prefer a single cubic over a fragmented fit when the chain is "almost"
# cubic: accept a single cubic if its max deviation is within this factor
# of tol. Keeps smooth curves as ONE editable cubic instead of dozens of
# fragments; genuine S-curves still split (their single-cubic error is far
# larger). The actual deviation is always reported per path.
SINGLE_CUBIC_TOL_FACTOR = 4.0
# Pre-fit smoothing sigma (upscaled px). Removes rasterization staircase
# noise (high frequency) without distorting real curve shape (which bends
# over tens of px). Applied `smooth_iters` times per SMOOTH_PRESETS.
SMOOTH_SIGMA = 1.5
class VectorizerConfig:
threshold: int = 110
upscale: int = 3 # Quality menu: 1x/2x/3x/4x
smoothing: str = "Low" # Low / Medium / High
tolerance: float = 0.6 # px, upscaled space
detect_corners: bool = True
detect_junctions: bool = True
detect_circles: bool = True # circles AND arcs
reduce_anchors: bool = False
use_sam: bool = False # MobileSAM refinement (slow, optional)
use_clip: bool = True # CLIP semantic labels (optional)