""" mosketch_lib.py — all non-Qwen, non-CLI logic for the sketch deformation pipeline in one place: loading SVG/semantic/trajectory/caption files, building meshes, running ARAP, and rendering results. This file has NO command-line interface and makes NO calls to Qwen or any other model — it's a pure library, imported by mosketch_pipeline.py. Consolidated from (all logic ported, not rewritten): arap_2d.py, trajectory.py, captions.py, and the mesh/handle-selection functions from auto_pipeline.py (auto_select_handles_deduped and its dependencies — auto_pipeline.py's own moving/static classifier and rigidity heuristic were NOT ported here, since both were confirmed unused by the real pipeline and the rigidity heuristic was confirmed broken on real data). """ import json import os import re import numpy as np import scipy.sparse as sp import scipy.sparse.linalg as spla from scipy.spatial import Delaunay, cKDTree from sklearn.cluster import KMeans, DBSCAN from svgpathtools import svg2paths # ============================================================================= # SVG / semantic loading # ============================================================================= def load_strokes_from_svg(svg_path, samples_per_stroke=12): paths, _ = svg2paths(svg_path) strokes = [] for p in paths: if len(p) == 0: continue ts = np.linspace(0, 1, samples_per_stroke) pts = np.array([[p.point(t).real, p.point(t).imag] for t in ts]) strokes.append(pts) if not strokes: raise ValueError(f"No usable paths found in {svg_path}") return strokes def load_semantic_assignments(path): """Returns dict: object_name -> list of int stroke indices.""" assignments = {} with open(path) as f: for line in f: line = line.rstrip("\n") if not line.strip(): continue name, idx_str = line.split("\t") indices = [int(i) for i in idx_str.split(",") if i.strip() != ""] assignments[name] = indices return assignments def filter_strokes(strokes, indices): return [strokes[i] for i in indices] def flatten_strokes(strokes): """Returns (all_points (N,2), stroke_slices [(start,end), ...]).""" all_points = [] stroke_slices = [] offset = 0 for stroke in strokes: n = len(stroke) stroke_slices.append((offset, offset + n)) all_points.append(stroke) offset += n return np.vstack(all_points), stroke_slices def load_object(name, svg_path, semantic_path): """Convenience: load one named object's flattened points + slices directly.""" strokes = load_strokes_from_svg(svg_path) assignments = load_semantic_assignments(semantic_path) strokes = filter_strokes(strokes, assignments[name]) return flatten_strokes(strokes) # ============================================================================= # Mesh construction (dedup + Delaunay) # ============================================================================= def deduplicate_points(points, tol=1e-3): """ Merges points within `tol` of each other into a single mesh vertex. Necessary because strokes meeting at a shared joint often have exactly (or near-exactly) coincident coordinates — Delaunay/QHull can leave one copy of a duplicate pair disconnected from every triangle, which makes the ARAP Laplacian singular. Returns: unique_points: (M, 2) deduplicated point positions. point_to_unique: (N,) int array mapping each original point index to its merged vertex index in unique_points. """ N = len(points) tree = cKDTree(points) parent = list(range(N)) def find(x): while parent[x] != x: parent[x] = parent[parent[x]] x = parent[x] return x def union(a, b): ra, rb = find(a), find(b) if ra != rb: parent[ra] = rb pairs = tree.query_pairs(r=tol) for a, b in pairs: union(a, b) roots = sorted(set(find(i) for i in range(N))) root_to_unique = {r: i for i, r in enumerate(roots)} point_to_unique = np.array([root_to_unique[find(i)] for i in range(N)]) unique_points = np.zeros((len(roots), 2)) counts = np.zeros(len(roots)) for i in range(N): u = point_to_unique[i] unique_points[u] += points[i] counts[u] += 1 unique_points /= counts[:, None] return unique_points, point_to_unique def build_mesh(points): """Delaunay-triangulates the (already deduplicated) point cloud.""" tri = Delaunay(points) edges = set() for simplex in tri.simplices: for i in range(3): a, b = simplex[i], simplex[(i + 1) % 3] edges.add((min(a, b), max(a, b))) return tri, np.array(sorted(edges)) def nearest_mesh_vertex(unique_points, target): dists = np.linalg.norm(unique_points - np.array(target), axis=1) return int(dists.argmin()) # ============================================================================= # Automatic joint/handle selection (KMeans + collision-safe dedup) # ============================================================================= def auto_select_handles(strokes, k=4, seed=0): """ Runs KMeans on stroke centroids to get k joints. The joint nearest the object's own centroid becomes the anchor; all others become handles. Returns RAW (non-deduplicated) joint positions — use auto_select_handles_deduped for the mesh-safe version. """ centroids = np.array([s.mean(axis=0) for s in strokes]) km = KMeans(n_clusters=k, n_init=10, random_state=seed).fit(centroids) joints = km.cluster_centers_ object_centroid = centroids.mean(axis=0) dists = np.linalg.norm(joints - object_centroid, axis=1) anchor_idx = int(dists.argmin()) handle_idxs = [i for i in range(k) if i != anchor_idx] return joints, anchor_idx, handle_idxs def cluster_object_strokes(points, slices, bbox_size, eps_fraction=0.08, min_samples=2): """ Groups an object's own strokes into spatially distinct clusters via DBSCAN on each stroke's centroid — a purely geometric candidate-region detector, no VLM involved. WHY THIS EXISTS: adapted from the reference paper's approach (cone singularity detection computes candidate handle points BEFORE the VLM ever sees them; the VLM only ever selects among pre-computed real candidates, never invents a coordinate from scratch). That paper's exact mechanism doesn't port directly — it relies on 3D mesh sub-part segmentation we don't have (see run_semantic_handle_selection's docstring for the full reasoning) — but the underlying principle generalizes: candidates should come from real geometry, not a free VLM guess that then gets silently accepted regardless of how far it drifts from anything real. CONFIRMED on real hardware (eat2): a VLM twice fabricated a second hand and once placed "chin" inside the hair, with coordinates that snapped to real-but-wrong mesh vertices — the free-coordinate approach has no way to prevent this, since nearest_mesh_vertex always returns SOMETHING. Tested this clustering directly against eat2's real stroke geometry: the hand+spoon tangle (34 strokes) forms exactly ONE clean, spatially isolated cluster, fully separate from the face+hair mass — meaning "claim there are two hands" becomes structurally impossible if the VLM can only select among real candidate clusters, since there is only one hand-region cluster to point at. eps_fraction: DBSCAN's eps (max distance between two stroke centroids to be considered the same cluster) as a fraction of the object's own bbox_size — scales with sketch/canvas size rather than being a fixed absolute pixel value, same convention as HANDLE_MOVE_CAP_FRACTION elsewhere in this codebase. 0.08 matched the real eat2 test (eps=20px at bbox_size=238.4). min_samples=2: a cluster needs at least 2 strokes to count as a real region — a single stray stroke (common in a rough sketch) becomes DBSCAN noise (label -1) rather than its own spurious one-stroke "part", since a real body part is normally drawn with more than one stroke. Returns a list of dicts: [{"cluster_id": int, "centroid": [x,y], "n_strokes": int, "stroke_indices": [...]}, ...] — sorted by cluster_id, noise strokes (label -1) excluded entirely (not returned as a candidate — they're not a coherent region to name). """ if len(slices) == 0: return [] stroke_centroids = np.array([points[start:end].mean(axis=0) for start, end in slices]) eps = max(bbox_size * eps_fraction, 1.0) # floor of 1px guards against a degenerate bbox_size labels = DBSCAN(eps=eps, min_samples=min_samples).fit(stroke_centroids).labels_ clusters = [] for cluster_id in sorted(set(labels)): if cluster_id == -1: continue # noise — not a coherent candidate region stroke_indices = [i for i, lbl in enumerate(labels) if lbl == cluster_id] centroid = stroke_centroids[stroke_indices].mean(axis=0) clusters.append({ "cluster_id": int(cluster_id), "centroid": [float(centroid[0]), float(centroid[1])], "n_strokes": len(stroke_indices), "stroke_indices": stroke_indices, }) return clusters def auto_select_handles_deduped(strokes, unique_points, k=4, seed=0): """ Like auto_select_handles, but maps each joint to its nearest mesh vertex first and DROPS any joint that collides with an already-claimed vertex. Necessary for small objects (e.g. a frisbee) where KMeans can find k "distinct" cluster centers that all round to the same underlying mesh point after deduplication — without this, two different joint names would silently control the same physical vertex, and whichever target arrives last would overwrite the other with no error. Returns: joints: (M, 2) deduped joint positions, M <= k. anchor_idx: index into `joints`. handle_idxs: list of indices into `joints`. joint_mesh_indices: list of mesh vertex indices, one per joint, guaranteed unique. """ joints_raw, _, _ = auto_select_handles(strokes, k=k, seed=seed) seen_vertices = {} kept_joints = [] kept_mesh_indices = [] for j in joints_raw: mv = nearest_mesh_vertex(unique_points, j) if mv in seen_vertices: continue seen_vertices[mv] = True kept_joints.append(j) kept_mesh_indices.append(mv) kept_joints = np.array(kept_joints) object_centroid = kept_joints.mean(axis=0) dists = np.linalg.norm(kept_joints - object_centroid, axis=1) anchor_idx = int(dists.argmin()) handle_idxs = [i for i in range(len(kept_joints)) if i != anchor_idx] dropped = k - len(kept_joints) if dropped > 0: print(f" NOTE: {dropped} of {k} joints collided onto an already-used mesh vertex " f"and were dropped — object only supports {len(kept_joints)} independent handles") return kept_joints, anchor_idx, handle_idxs, kept_mesh_indices def object_bbox_size(points): return np.linalg.norm(points.max(axis=0) - points.min(axis=0)) def build_stroke_geometry_text(strokes, n_points=2): """ Formats an object's actual stroke points as compact text — e.g. "stroke_0: (12,45) -> (22,53)" per stroke. Built specifically to give Qwen the object's real geometry in TEXT form, not just a rendered image: there's real literature suggesting multimodal LLMs can under-attend to image tokens relative to text, so this is a complementary, text-native representation of the same information the image already shows, not a replacement for it. n_points=2 (start+end of each stroke only) was chosen after measuring real cost on dog3's 64-stroke dog object: 2 points/stroke = ~570 tokens, 4 points/stroke = ~990 tokens, full 12 points/stroke = ~2650 tokens — given this pipeline has already hit real OOM/latency limits this session, the sparsest still-useful setting was chosen deliberately, not the richest one. """ lines = [] for i, stroke in enumerate(strokes): idxs = np.linspace(0, len(stroke) - 1, n_points).astype(int) pts = stroke[idxs] pts_str = " -> ".join(f"({x:.0f},{y:.0f})" for x, y in pts) lines.append(f" stroke_{i}: {pts_str}") return "\n".join(lines) # ============================================================================= # ARAP solver (Sorkine & Alexa 2007, local-global, uniform edge weights) # ============================================================================= def arap_deform(points, edges, handle_indices, handle_targets, iterations=10): """ points: (N, 2) rest-pose positions. edges: (E, 2) int array of vertex index pairs. handle_indices: list of point indices with fixed target positions. handle_targets: (len(handle_indices), 2) target positions. Returns: deformed points, (N, 2). """ N = len(points) neighbors = {i: [] for i in range(N)} for a, b in edges: neighbors[a].append(b) neighbors[b].append(a) W = sp.lil_matrix((N, N)) for a, b in edges: W[a, b] = 1.0 W[b, a] = 1.0 W = W.tocsr() deg = np.array(W.sum(axis=1)).flatten() L = sp.diags(deg) - W p_prime = points.copy() for idx, target in zip(handle_indices, handle_targets): p_prime[idx] = target L_constrained = L.tolil() for idx in handle_indices: L_constrained.rows[idx] = [idx] L_constrained.data[idx] = [1.0] L_constrained = L_constrained.tocsr() solver = spla.factorized(L_constrained.tocsc()) for it in range(iterations): rotations = np.zeros((N, 2, 2)) for i in range(N): js = neighbors[i] if not js: rotations[i] = np.eye(2) continue P = np.array([points[i] - points[j] for j in js]).T Pp = np.array([p_prime[i] - p_prime[j] for j in js]).T S = P @ Pp.T U, _, Vt = np.linalg.svd(S) R = Vt.T @ U.T if np.linalg.det(R) < 0: Vt[-1, :] *= -1 R = Vt.T @ U.T rotations[i] = R bx = np.zeros(N) by = np.zeros(N) for a, b in edges: e_rest = points[a] - points[b] contrib = 0.5 * (rotations[a] + rotations[b]) @ e_rest bx[a] += contrib[0]; by[a] += contrib[1] bx[b] -= contrib[0]; by[b] -= contrib[1] for idx, target in zip(handle_indices, handle_targets): bx[idx] = target[0] by[idx] = target[1] x_new = solver(bx) y_new = solver(by) p_prime = np.stack([x_new, y_new], axis=1) return p_prime # ============================================================================= # Trajectory parsing (real _traj.txt files) # ============================================================================= DEFAULT_KEYFRAME_INDICES = (0, 4, 9, 14, 19) def load_trajectories(traj_path, keyframe_indices=DEFAULT_KEYFRAME_INDICES): """ File format: tab-separated, one object per line — object_name\\t[x,y,w,h],[x,y,w,h],... (20 frames) Returns dict: object_name -> list of (x,y,w,h), sampled at keyframe_indices. """ trajectories = {} with open(traj_path) as f: for line_no, line in enumerate(f, 1): line = line.rstrip("\n") if not line.strip(): continue if "\t" not in line: raise ValueError( f"{traj_path} line {line_no}: expected tab-separated " f"'object_name\\t[x,y,w,h],...' format, got: {line[:80]!r}" ) name, boxes_str = line.split("\t", 1) box_strs = re.findall(r"\[([^\]]+)\]", boxes_str) if not box_strs: raise ValueError(f"{traj_path} line {line_no}: no [x,y,w,h] boxes found for '{name}'") boxes = [tuple(float(v) for v in b.split(",")) for b in box_strs] max_idx = max(keyframe_indices) if len(boxes) <= max_idx: raise ValueError( f"{traj_path} line {line_no}: object '{name}' has only {len(boxes)} " f"frames, need at least {max_idx + 1} for keyframe indices {keyframe_indices}" ) trajectories[name] = [boxes[i] for i in keyframe_indices] return trajectories def bbox_deltas(sampled_boxes): """Returns (dx, dy) numpy arrays, offset from the first sampled keyframe.""" xs = np.array([b[0] for b in sampled_boxes], dtype=float) ys = np.array([b[1] for b in sampled_boxes], dtype=float) return xs - xs[0], ys - ys[0] # ============================================================================= # Caption parsing (real caption.txt files) # ============================================================================= def load_captions(caption_path): """ File format (CONFIRMED against the real 60-sketch caption.txt): tab-separated, one sketch per line — sketch_name\\tcaption text. Returns dict: sketch_name -> caption string. """ captions = {} with open(caption_path) as f: for line_no, line in enumerate(f, 1): line = line.rstrip("\n") if not line.strip(): continue if "\t" not in line: raise ValueError( f"{caption_path} line {line_no}: expected tab-separated " f"'sketch_name\\tcaption' format, got: {line[:80]!r}" ) name, caption = line.split("\t", 1) captions[name.strip()] = caption.strip() return captions def get_caption(caption_path, sketch_name): captions = load_captions(caption_path) if sketch_name not in captions: raise KeyError( f"'{sketch_name}' not found in {caption_path}. " f"Available: {list(captions.keys())[:10]}{'...' if len(captions) > 10 else ''}" ) return captions[sketch_name] # ============================================================================= # Structured interaction constraints (per-sketch, optional) # ============================================================================= def repair_json_trailing_commas(text): """ Strips a trailing comma immediately before a closing '}' or ']' (across whitespace/ newlines) — the single most common way small VLMs produce almost-valid JSON: a syntactically-correct-looking object/list with one extra comma before the close. CONFIRMED on real hardware (eat2, judge attempt 3): a verdict with a trailing comma after the last motion_checks entry — "temporal_coherence": {...},\n}\n``` — failed json.loads with "Expecting property name enclosed in double quotes", and the entire attempt (a full judge call, real tokens, real GPU time) was discarded and retried from scratch over one stray comma, even though every other field in that response was well-formed and usable. Deliberately narrow: only removes a comma that is followed by nothing but whitespace and a closing bracket. This does NOT attempt to repair other malformed JSON (unquoted keys, single quotes, unescaped strings, genuinely missing brackets) — those are different failure modes that a targeted trailing-comma strip could accidentally paper over if it tried to be more permissive, and a response that's broken in one of THOSE ways should still fail loudly rather than be guessed into some other, potentially wrong, shape. Safe to call unconditionally before json.loads — running it on already-valid JSON is a no-op (there is nothing for the pattern to match). """ return re.sub(r",(\s*[}\]])", r"\1", text) def load_interaction_constraints(path): """ Loads structured object-to-object interaction constraints for one sketch, e.g.: [ {"source_object": "dog", "source_part": "mouth", "target_object": "frisbee", "relationship": "mouth near target", "critical_keyframe": 4} ] This is the "structured interaction constraints" item from the handoff document's TODO list (the dog/mouth/frisbee example is theirs, kept verbatim as the canonical shape). It is NOT derived automatically from the caption — captions are prose, not structured data, and reliably turning them into (source_part, target_object, keyframe) triples is its own unsolved problem, not something this change attempts. Instead this is a small, optionally-authored file per sketch, same pattern as {name}_semantic.txt. File is OPTIONAL: if `path` doesn't exist, returns [] and callers fall back to the old behavior (no interaction-specific checks). This keeps every sketch without a file working exactly as before — nothing about this is required. Each entry is validated minimally on load (required keys present, critical_keyframe is an int) — malformed entries are dropped with a printed warning rather than raising, since a typo in a hand-authored file shouldn't crash a whole pipeline run. """ if not os.path.exists(path): return [] with open(path) as f: raw = json.load(f) if not isinstance(raw, list): raise ValueError(f"{path}: expected a JSON list of interaction constraints, got {type(raw).__name__}") required = ("source_object", "source_part", "target_object", "relationship", "critical_keyframe") constraints = [] for i, entry in enumerate(raw): if not isinstance(entry, dict): print(f"WARNING: {path}[{i}] is not an object, skipping: {entry!r}") continue missing = [k for k in required if k not in entry] if missing: print(f"WARNING: {path}[{i}] missing keys {missing}, skipping: {entry!r}") continue if not isinstance(entry["critical_keyframe"], int): print(f"WARNING: {path}[{i}]['critical_keyframe'] = {entry['critical_keyframe']!r} is not an " f"int, skipping") continue constraints.append({k: entry[k] for k in required}) return constraints # ============================================================================= # Rendering # ============================================================================= def render(points, stroke_slices, mesh_edges=None, handles=None, out_path="out.png", title=""): import matplotlib.pyplot as plt fig, ax = plt.subplots(figsize=(6, 6)) if mesh_edges is not None: for a, b in mesh_edges: ax.plot([points[a, 0], points[b, 0]], [points[a, 1], points[b, 1]], color="#D9DEE2", linewidth=0.5, zorder=1) for start, end in stroke_slices: seg = points[start:end] ax.plot(seg[:, 0], seg[:, 1], color="black", linewidth=1.3, zorder=2) if handles is not None: ax.scatter(points[handles, 0], points[handles, 1], color="#3D7A5C", s=80, zorder=5) ax.invert_yaxis() ax.set_aspect("equal") ax.set_title(title) fig.savefig(out_path, dpi=150, bbox_inches="tight") plt.close(fig)