diff --git a/.gitignore b/.gitignore index a3b109b..869283f 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ .idea/ *.egg-info /build/ +/bench/results/ /output /models # Run output: vectrify writes /runs/ next to the output file diff --git a/README.md b/README.md index 93c2a0c..0516979 100644 --- a/README.md +++ b/README.md @@ -61,8 +61,8 @@ vectrify photo.jpg -o sketch.svg --seeds 10 --epochs 4 \ --max-wall-seconds 1800 vectrify mascot.png -o mascot.svg --segment-count 12 # tiles/local elites (default: 8) -# Skip the default segmentation-derived SVG seed (SVG only) -vectrify artwork.png -o artwork.svg --no-samvg-seed +# Add the optional segmentation-derived SVG seed (SVG only) +vectrify artwork.png -o artwork.svg --samvg-seed # Choose a provider, model, or scorer explicitly vectrify input.png --provider anthropic --model MODEL_NAME @@ -109,14 +109,13 @@ by default; add `--save-heatmap` for perceptual difference maps. ## SAMVG-inspired seed -One native, segmentation-first SVG candidate is added to every SVG run without -reducing the configured LLM seed count. It uses SAM ViT-H by default, retains -masks only when they materially improve a flat-colour reconstruction of the -target, and traces the retained masks into editable layered SVG paths. Set +With `--samvg-seed`, Vectrify adds one native, segmentation-first SVG candidate +without reducing the configured LLM seed count. It uses SAM ViT-H by default, +retains masks only when they materially improve a flat-colour reconstruction of +the target, and traces the retained masks into editable layered SVG paths. Set `VECTRIFY_SAMVG_MODEL=facebook/sam-vit-base` for the smaller checkpoint. It is -inspired by SAMVG, not an installation of the unreleased research code. Use -`--no-samvg-seed` to skip it; the feature is currently available for SVG output -only. +inspired by SAMVG, not an installation of the unreleased research code and is +off by default. SAM inputs default to a 1024px maximum side, the model's native encoder size; the returned masks are restored to the target's original canvas before tracing. diff --git a/scripts/bench_samvg_two_phase.py b/scripts/bench_samvg_two_phase.py index 69bb8c4..614fdc4 100644 --- a/scripts/bench_samvg_two_phase.py +++ b/scripts/bench_samvg_two_phase.py @@ -28,6 +28,7 @@ _mse, _render_layers, _render_svg, + _sam_runtime, filter_by_impact, prompted_masks, residual_prompt_points, @@ -43,13 +44,11 @@ def _path_count(svg: str) -> int: ) -def _l1(target: Image.Image, rendered: Image.Image) -> float: - return float( - np.abs( - np.asarray(target.convert("RGB"), dtype=np.float32) / 255.0 - - np.asarray(rendered.convert("RGB"), dtype=np.float32) / 255.0 - ).mean() - ) +def _result_name(target_path: Path) -> str: + """Give standard bench targets stable, non-colliding output directories.""" + if target_path.name == "target.png": + return target_path.parent.name + return target_path.stem def _write_gallery(images: list[tuple[str, Image.Image]], destination: Path) -> None: @@ -68,6 +67,7 @@ def _fit_if_improved( target: Image.Image, plugin: SvgPlugin, steps: int, + learn_alpha: bool, ) -> tuple[str, Image.Image, list[dict[str, int | float]], bool]: before = _render_svg(svg, target, plugin.rasterize) measurements: list[dict[str, int | float]] = [] @@ -77,6 +77,7 @@ def _fit_if_improved( rasterize=plugin.rasterize, steps=steps, measurements=measurements, + learn_alpha=learn_alpha, ) after = _render_svg(candidate, target, plugin.rasterize) if _mse(target, after) <= _mse(target, before): @@ -90,36 +91,101 @@ def run_target( *, steps: int, reference_svg: Path | None = None, + learn_alpha: bool = False, + curvature_threshold: float | None = None, + seed_only: bool = False, ) -> None: target = Image.open(target_path).convert("RGB") plugin = SvgPlugin() - destination = output / target_path.stem + destination = output / _result_name(target_path) destination.mkdir(parents=True, exist_ok=True) started = perf_counter() - layers = retrieve_layers(target) + runtime = _sam_runtime() + layers = retrieve_layers(target, _runtime=runtime) initial = _append_layers( f'', layers, 16, hybrid_strokes=False, + curvature_threshold=curvature_threshold, ) + _render_svg(initial, target, plugin.rasterize).save(destination / "first-seed.png") + (destination / "first-seed.svg").write_text(initial) + if seed_only: + seed_render = _render_svg(initial, target, plugin.rasterize) + mask_canvas, _coverage = _render_layers((target.height, target.width), layers) + stages = [ + ("target", target, None), + ("mask-canvas", Image.fromarray(mask_canvas), None), + ("first-seed", seed_render, initial), + ] + if reference_svg is not None: + reference = _render_svg(reference_svg.read_text(), target, plugin.rasterize) + stages.append(("reference-svg", reference, reference_svg.read_text())) + rows = [ + { + "stage": name, + "mse": _mse(target, rendered), + "paths": _path_count(svg) if svg is not None else 0, + } + for name, rendered, svg in stages + ] + _write_gallery( + [(name, image) for name, image, _svg in stages], destination / "gallery.png" + ) + with (destination / "stages.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=["stage", "mse", "paths"]) + writer.writeheader() + writer.writerows(rows) + (destination / "fit-groups.json").write_text("[]\n") + (destination / "summary.json").write_text( + json.dumps( + { + "target": str(target_path), + "seed_only": True, + "initial_layers": len(layers), + "wall_seconds": perf_counter() - started, + "stages": rows, + }, + indent=2, + ) + ) + return first, first_render, first_measurements, first_accepted = _fit_if_improved( - initial, target, plugin, steps + initial, target, plugin, steps, learn_alpha ) + first_render.save(destination / "first-fit.png") + (destination / "first-fit.svg").write_text(first) points = residual_prompt_points(target, first_render) - _canvas, coverage = _render_layers((target.height, target.width), layers) added = filter_by_impact( target, - prompted_masks(target, points), + prompted_masks(target, points, _runtime=runtime), existing=layers, initial_canvas=np.asarray(first_render, dtype=np.uint8), - initial_coverage=coverage, + # The residual pass starts from the first fitted raster. It is not an + # uncovered-mask pass, so all pixels must use their actual raster MSE. + initial_coverage=np.ones((target.height, target.width), dtype=bool), )[len(layers) :] - recovery = _append_layers(first, added, 16, hybrid_strokes=False) + recovery = _append_layers( + first, + added, + 16, + hybrid_strokes=False, + curvature_threshold=curvature_threshold, + ) + _render_svg(recovery, target, plugin.rasterize).save( + destination / "residual-recovery.png" + ) + (destination / "residual-recovery.svg").write_text(recovery) final, final_render, final_measurements, final_accepted = _fit_if_improved( - recovery, target, plugin, steps + recovery, target, plugin, steps, learn_alpha ) + if _mse(target, final_render) > _mse(target, first_render): + # Phase-two fitting is accepted relative to the recovered document, + # but the benchmark's final result must retain the already accepted + # first fit when residual additions regress the exported Cairo raster. + final, final_render, final_accepted = first, first_render, False stages = [ ("target", target, None), ("first-seed", _render_svg(initial, target, plugin.rasterize), initial), @@ -142,7 +208,6 @@ def run_target( rows.append( { "stage": name, - "l1": _l1(target, rendered), "mse": _mse(target, rendered), "paths": _path_count(svg) if svg is not None else 0, } @@ -151,7 +216,7 @@ def run_target( [(name, image) for name, image, _svg in stages], destination / "gallery.png" ) with (destination / "stages.csv").open("w", newline="") as handle: - writer = csv.DictWriter(handle, fieldnames=["stage", "l1", "mse", "paths"]) + writer = csv.DictWriter(handle, fieldnames=["stage", "mse", "paths"]) writer.writeheader() writer.writerows(rows) measurements = [ @@ -182,6 +247,11 @@ def main() -> None: type=Path, help="Reference SVG to Cairo-rasterize alongside a single target.", ) + parser.add_argument( + "--curvature-threshold", + type=float, + help="Use the dissertation's variable-segment tracing variation.", + ) parser.add_argument("--cat", action="store_true") parser.add_argument( "--all", action="store_true", help="Run cat, duck, and all bench targets." @@ -190,6 +260,16 @@ def main() -> None: "--output", type=Path, default=ROOT / "bench/results/samvg-two-phase" ) parser.add_argument("--steps", type=int, default=500) + parser.add_argument( + "--seed-only", + action="store_true", + help="Benchmark automatic masks and coverage recovery without fitting.", + ) + parser.add_argument( + "--learn-alpha", + action="store_true", + help="Use the dissertation's SAMVG+alpha fitter variation.", + ) args = parser.parse_args() targets = list(args.target) if args.cat or args.all: @@ -214,6 +294,9 @@ def main() -> None: args.output, steps=args.steps, reference_svg=reference_svg, + learn_alpha=args.learn_alpha, + curvature_threshold=args.curvature_threshold, + seed_only=args.seed_only, ) diff --git a/src/vectrify/cli.py b/src/vectrify/cli.py index b61c0fd..50e1c58 100644 --- a/src/vectrify/cli.py +++ b/src/vectrify/cli.py @@ -4,6 +4,11 @@ from importlib.metadata import version as _pkg_version from vectrify.formats import FORMAT_NAMES +from vectrify.refine.samvg import ( + SAMVG_MAX_SIDE, + SAMVG_MODEL, + SAMVG_POINTS_PER_BATCH, +) from vectrify.score import ScorerType from vectrify.score.vision import DEFAULT_VISION_MODEL @@ -196,14 +201,87 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: metavar="N", help="Edge-aware Voronoi masks retained as local elites. Default: 8", ) - g_search.add_argument( + g_samvg = parser.add_argument_group("SAMVG seed") + g_samvg.add_argument( "--samvg-seed", action=argparse.BooleanOptionalAction, - default=True, + default=False, help="Add one SAMVG-inspired SVG seed made from automatic SAM masks, " "impact filtering, contour tracing, and Torch OCR. Requires " "vectrify[samvg] and " - "is available for SVG output only. Default: on", + "is available for SVG output only. Default: off", + ) + g_samvg.add_argument( + "--samvg-model", + default=SAMVG_MODEL, + metavar="HF_REPO", + help=f"HuggingFace SAM checkpoint used for the seed. Default: {SAMVG_MODEL}", + ) + g_samvg.add_argument( + "--samvg-max-side", + type=int, + default=SAMVG_MAX_SIDE, + metavar="PX", + help="Maximum long side passed to SAM before masks are restored to " + "the output canvas. " + f"Default: {SAMVG_MAX_SIDE}", + ) + g_samvg.add_argument( + "--samvg-points-per-batch", + type=int, + default=SAMVG_POINTS_PER_BATCH, + metavar="N", + help="SAM decoder prompts per GPU batch; this does not change the " + "32x32 automatic prompt grid. " + f"Default: {SAMVG_POINTS_PER_BATCH}", + ) + g_samvg.add_argument( + "--samvg-min-pixels", + type=int, + default=32, + metavar="N", + help="Discard connected mask components smaller than N pixels before " + "tracing. Default: 32", + ) + g_samvg.add_argument( + "--samvg-min-impact", + type=float, + default=3e-6, + metavar="MSE", + help="Minimum whole-mask reconstruction improvement required for " + "retention. Default: 3e-6", + ) + g_samvg.add_argument( + "--samvg-max-layers", + type=int, + default=512, + metavar="N", + help="Maximum retained SAM masks in each seed stage. Default: 512", + ) + g_samvg.add_argument( + "--samvg-segments", + type=int, + default=16, + metavar="N", + help="Fixed cubic Bézier segments per traced contour. Default: 16", + ) + g_samvg.add_argument( + "--samvg-fill-holes", + action=argparse.BooleanOptionalAction, + default=True, + help="Fill only sub-threshold enclosed mask holes before tracing. Default: on", + ) + g_samvg.add_argument( + "--samvg-hybrid-strokes", + action=argparse.BooleanOptionalAction, + default=True, + help="Emit conservative centreline strokes for thin seed masks. Default: on", + ) + g_samvg.add_argument( + "--samvg-ocr", + action=argparse.BooleanOptionalAction, + default=True, + help="Run optional OCR and retain pixel-verified editable text. Default: on", ) g_epoch = parser.add_argument_group( @@ -484,6 +562,16 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace: raise SystemExit("Error: --workers and --pool-size must be > 0") if ns.segment_count <= 0: raise SystemExit("Error: --segment-count must be > 0") + if ( + ns.samvg_max_side <= 0 + or ns.samvg_points_per_batch <= 0 + or ns.samvg_min_pixels <= 0 + or ns.samvg_max_layers <= 0 + or ns.samvg_segments <= 0 + ): + raise SystemExit("Error: SAMVG integer controls must be > 0") + if ns.samvg_min_impact < 0: + raise SystemExit("Error: --samvg-min-impact must be >= 0") if ns.resolution <= 0: raise SystemExit("Error: --resolution must be > 0") if ns.resolution_llm <= 0: diff --git a/src/vectrify/main.py b/src/vectrify/main.py index 7e6c33d..48a102e 100755 --- a/src/vectrify/main.py +++ b/src/vectrify/main.py @@ -153,6 +153,16 @@ def main(): vision_model=args.vision_model, segment_count=args.segment_count, samvg_seed=args.samvg_seed, + samvg_model=args.samvg_model, + samvg_max_side=args.samvg_max_side, + samvg_points_per_batch=args.samvg_points_per_batch, + samvg_min_pixels=args.samvg_min_pixels, + samvg_min_impact=args.samvg_min_impact, + samvg_max_layers=args.samvg_max_layers, + samvg_segments=args.samvg_segments, + samvg_fill_holes=args.samvg_fill_holes, + samvg_hybrid_strokes=args.samvg_hybrid_strokes, + samvg_ocr=args.samvg_ocr, auto_crop=args.auto_crop, dry_run=args.dry_run, dry_run_parameters=vars(args), diff --git a/src/vectrify/refine/paths.py b/src/vectrify/refine/paths.py index 4f2d16e..812c7ac 100644 --- a/src/vectrify/refine/paths.py +++ b/src/vectrify/refine/paths.py @@ -11,6 +11,7 @@ import itertools import logging import math +import os import random import re from collections import defaultdict @@ -580,6 +581,15 @@ def _dynamic_fill_winding_chunk(start: Any, end: Any, pixels: Any) -> Any: return _fill_winding_chunk(start, end, pixels) +def _torch_compile_enabled() -> bool: + """Whether this process permits Torch Dynamo to compile renderer kernels.""" + # PyTorch raises from an already-created ``torch.compile`` wrapper when + # this environment switch is set. Respect it before creating the cached + # wrapper so the advertised eager renderer is actually usable for + # debugging, constrained deployments, and compiler-cache recovery. + return os.environ.get("TORCH_COMPILE_DISABLE", "0") not in {"1", "true", "True"} + + @lru_cache(maxsize=1) def _compiled_fill_winding_chunk() -> Any: """Return the CUDA-fused winding primitive when this torch supports it. @@ -592,7 +602,7 @@ def _compiled_fill_winding_chunk() -> Any: import torch compile_fn = getattr(torch, "compile", None) - if compile_fn is None: + if compile_fn is None or not _torch_compile_enabled(): return _fill_winding_chunk try: return compile_fn( @@ -615,7 +625,7 @@ def _compiled_tiled_fill_winding_chunk() -> Any: import torch compile_fn = getattr(torch, "compile", None) - if compile_fn is None: + if compile_fn is None or not _torch_compile_enabled(): return _fill_winding_chunk try: return compile_fn( @@ -661,21 +671,30 @@ def _fill_batched_windings( """ import torch - # The packaged native operator uses SAMVG's fixed 16-cubic contour - # representation. It remains deliberately narrow: every - # other contour/layout continues through the proven Torch implementation. + # The native primitive is fixed-width, but winding is additive over cubic + # ranges. Chunk a long contour into padded 16-cubic ranges and sum its + # exact native winding fields before applying the SVG fill rule. This is + # the same representation as the fixed SAMVG path, not a tessellation or + # geometry approximation, and keeps SAMVG+var off Torch's huge broadcast + # fallback. if samples in {8, 16, 32}: from vectrify.refine.cuda_renderer import winding as cuda_winding + chunks = math.ceil(controls.shape[1] / _FUSED_CUBICS) + padded = controls + if chunks > 1: + count = chunks * _FUSED_CUBICS - controls.shape[1] + point = controls[:, :1, :1].expand(-1, count, 4, -1) + padded = torch.cat((controls, point), dim=1) native = cuda_winding( - controls, + padded.reshape(-1, _FUSED_CUBICS, 4, 2), box, samples=samples, x_offset=x_offset, y_offset=y_offset, ) if native is not None: - return native + return native.reshape(len(controls), chunks, *native.shape[1:]).sum(dim=1) left, top, right, bottom = box height, width = bottom - top, right - left @@ -1365,11 +1384,16 @@ def fit_filled_svg( monolithic: bool | None = None, curve_samples: int | None = None, backdrop: Image.Image | None = None, + learn_alpha: bool = False, + sparse_replay: bool = False, ) -> str: - """Optimise opaque filled cubic SVG paths against an RGB target. + """Optimise filled cubic SVG paths against an RGB target. SAMVG optimises opaque path coordinates and fill colours for 500 Adam - iterations in each of its two passes. This implementation reuses + iterations in each of its two passes. ``learn_alpha`` enables the + dissertation's SAMVG+alpha variation: each selected path receives a + learnable fill opacity. The standard SAMVG configuration deliberately + keeps it disabled and treats fills as opaque. This implementation reuses Vectrify's torch renderer instead of requiring DiffVG, while retaining the dissertation's full-resolution Adam defaults: point LR 1, colour LR .01, and MSE plus .02 Xing loss. It uses DiffVG's standard 2x2 optimisation @@ -1380,12 +1404,25 @@ def fit_filled_svg( is available only as an explicit caller-selected preview mode. CUDA uses one monolithic compositor graph for 64px-or-smaller working canvases by default; larger canvases retain the memory-bounded replay. + ``sparse_replay`` retains the same painter-order MSE derivative while + saving layer state only within each path's raster tile; it makes a full + 1024px SAMVG phase practical without a monolithic alpha stack. """ import xml.etree.ElementTree as ET import torch root = ET.fromstring(svg) + + def opacity(element) -> float: + """Read the directly applied SVG fill opacity, clamped for Adam.""" + try: + fill_opacity = float(element.get("fill-opacity", "1")) + element_opacity = float(element.get("opacity", "1")) + except ValueError: + return 1.0 + return min(1.0, max(0.0, fill_opacity * element_opacity)) + entries = [] for element in root.iter(): if element.tag.split("}")[-1] != "path" or not element.get("d"): @@ -1400,7 +1437,7 @@ def fit_filled_svg( fill_rule = element.get("fill-rule", "nonzero").strip().lower() if fill_rule not in {"evenodd", "nonzero"}: continue - entries.append((element, contours, colour, fill_rule)) + entries.append((element, contours, colour, fill_rule, opacity(element))) if not entries: raise UnsupportedPathError("no opaque filled cubic paths to optimise") @@ -1436,17 +1473,31 @@ def fit_filled_svg( ) for contour in contours ] - for _element, contours, _colour, _fill_rule in entries + for _element, contours, _colour, _fill_rule, _opacity in entries ] # A detailed SAMVG seed has hundreds of contours. Keeping each one as a # separate Adam parameter turns one optimiser update into hundreds of tiny - # CUDA kernels. Store fixed-width contour slots in one parameter and use + # CUDA kernels. Store equal-width contour slots in one parameter and use # narrow views below, retaining every original contour length in the SVG - # and Xing terms. + # and Xing terms. The native coverage primitive itself uses 16-cubic + # chunks, but SAMVG+var legitimately emits longer contours; storage must + # therefore use the document maximum rather than that renderer chunk size. flat_controls = [control for path in initial_controls for control in path] contour_sizes = [len(control) for control in flat_controls] + storage_width = max(contour_sizes) + + def pad_storage_control(control: Any) -> Any: + if len(control) == storage_width: + return control + return torch.cat( + ( + control, + control[:1].expand(storage_width - len(control), -1, -1), + ) + ) + control_storage = torch.nn.Parameter( - torch.cat([_pad_fused_cubics(control[None]) for control in flat_controls]) + torch.stack([pad_storage_control(control) for control in flat_controls]) ) controls = [] path_storage_spans = [] @@ -1464,7 +1515,7 @@ def fit_filled_svg( # avoids launching Adam's tiny update kernels once per SVG layer. color_storage = torch.nn.Parameter( torch.tensor( - [colour for _element, _contours, colour, _fill_rule in entries], + [colour for _element, _contours, colour, _fill_rule, _opacity in entries], dtype=torch.float32, device=device, ) @@ -1488,28 +1539,41 @@ def fit_filled_svg( device=device, ) ) + alpha_values = ( + torch.nn.Parameter( + torch.tensor( + [entry[4] for entry in entries], + dtype=torch.float32, + device=device, + ) + ) + if learn_alpha + else None + ) point_optimizer = torch.optim.Adam( [control_storage], lr=point_learning_rate, fused=device == "cuda" ) colour_optimizer = torch.optim.Adam( - [color_storage], lr=color_learning_rate, fused=device == "cuda" - ) - # The dissertation averages Xing within each contour then sums contours. - # Keep that weighting while evaluating the 413 cat contours in one CUDA - # expression rather than launching one tiny graph for each. - xing_contour_weights = torch.cat( - [ - torch.full( - (len(control),), - 1 / len(control), - dtype=goal.dtype, - device=device, - ) - for path in controls - for control in path - ] + [color_storage, *([] if alpha_values is None else [alpha_values])], + lr=color_learning_rate, + fused=device == "cuda", ) + def close_contours() -> None: + """Restore the shared joins of every traced closed Bezier contour. + + SAMVG traces closed fixed-segment loops. The packed parameter storage + keeps their cubic endpoints as separate Adam values for efficient + rasterisation, so project them back to a continuous closed contour + after each update. Otherwise a subpixel gap becomes an extra implicit + SVG closing cubic on export, violating the fixed-segment invariant. + """ + with torch.no_grad(): + for path in controls: + for contour in path: + contour[1:, 0].copy_(contour[:-1, 3]) + contour[-1, 3].copy_(contour[0, 0]) + def tile_for(path: list[Any]) -> tuple[int, int, int, int]: """A fixed, antialiased raster tile covering a path's control hull.""" points = torch.cat([control.detach().reshape(-1, 2) for control in path]) @@ -1520,11 +1584,19 @@ def tile_for(path: list[Any]) -> tuple[int, int, int, int]: right = min(work_width, math.ceil(float(points[:, 0].max())) + 2) bottom = min(work_height, math.ceil(float(points[:, 1].max())) + 2) - # Bucket dimensions keep many unrelated small paths in the same CUDA - # batch. Shift a bucket at the canvas edge rather than clipping its - # protected coverage margin. - tile_width = min(work_width, 8 * math.ceil(max(right - left, 1) / 8)) - tile_height = min(work_height, 8 * math.ceil(max(bottom - top, 1) / 8)) + # Bucket dimensions keep unrelated paths in the same CUDA batch. A + # 32px bucket roughly halves the distinct sizes of the 1024px cat + # seed versus 8px buckets while adding only a small protected fringe + # to the right/bottom of each tile. The origin—and therefore every + # coverage sample belonging to the path—remains unchanged. Shift a + # bucket at the canvas edge rather than clipping its antialias margin. + tile_bucket = 32 + tile_width = min( + work_width, tile_bucket * math.ceil(max(right - left, 1) / tile_bucket) + ) + tile_height = min( + work_height, tile_bucket * math.ceil(max(bottom - top, 1) / tile_bucket) + ) left = min(left, work_width - tile_width) top = min(top, work_height - tile_height) return left, top, tile_width, tile_height @@ -1566,12 +1638,30 @@ def cropped_simple_groups() -> dict[ ].append((index, left, top)) return groups - def rasterise_simple( + def rasterise_simple_tiles( fill_rule: str, tile_width: int, tile_height: int, items: list[tuple[int, int, int]], - ) -> list[tuple[int, Any]]: + ) -> list[tuple[int, Any, int, int]]: + # SAMVG+var can emit a contour longer than the fixed-width coverage + # primitive. Route those through the chunked native winding path; + # packing them into the old batched coverage call would force eager + # Torch broadcasting over every cubic and pixel. + if controls[items[0][0]][0].shape[0] > _FUSED_CUBICS: + output = [] + for index, left, top in items: + offset = controls[index][0].new_tensor((left, top)) + alpha = _fill_path_coverage( + [controls[index][0] - offset], + (0, 0, tile_width, tile_height), + fill_rule=fill_rule, + samples=samples_for(tile_width, tile_height), + subpixels=subpixels, + fuse=False, + ) + output.append((index, alpha, left, top)) + return output translated = torch.stack( [ controls[index][0] - controls[index][0].new_tensor((left, top)) @@ -1585,13 +1675,44 @@ def rasterise_simple( samples=samples_for(tile_width, tile_height), subpixels=subpixels, fuse=False, - dynamic_fuse=len(items) >= 4, + # Sparse replay keeps this graph alive through the layer's + # backward pass. Torch's dynamic compiler can specialise one + # large tile batch into an unbounded graph here; eager chunking + # has the same coverage/gradient while retaining the documented + # tile-local memory bound. + dynamic_fuse=False, ) return [ - (index, restore_tile(alpha, left, top)) + (index, alpha, left, top) for (index, left, top), alpha in zip(items, rasterised, strict=True) ] + def rasterise_simple( + fill_rule: str, + tile_width: int, + tile_height: int, + items: list[tuple[int, int, int]], + ) -> list[tuple[int, Any]]: + return [ + (index, restore_tile(alpha, left, top)) + for index, alpha, left, top in rasterise_simple_tiles( + fill_rule, tile_width, tile_height, items + ) + ] + + def sparse_backward_batch_size(tile_width: int, tile_height: int) -> int: + """Bound the live tile-local autograd graph while filling the GPU. + + Sparse replay never needs a full-canvas alpha stack, but it does keep + the coverage graph for one backward batch alive. A fixed 16-path + batch underutilises CUDA for SAMVG's common 32--64px tiles, but a + recovery pass can create a much larger equal-tile group than the + initial seed. Retain the proven 16-path graph cap and apply the + tile-area budget beneath it. This bounds peak memory for every + document without changing the rendered image or its derivative. + """ + return max(1, min(16, (1 << 20) // max(1, tile_width * tile_height))) + def rasterise_multi(index: int, path: list[Any]) -> Any: # Large paths use fixed conservative candidate tiles. Every tile # sees all contours that can cross one of its horizontal rays, while @@ -1661,6 +1782,20 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: ) left, top, tile_width, tile_height = initial_multi_tiles[index] offset = path[0].new_tensor((left, top)) + from vectrify.refine.cuda_renderer import multi_coverage + + packed = torch.cat( + [_pad_fused_cubics((control - offset)[None]) for control in path] + ) + analytic = multi_coverage( + packed, + [0, len(path)], + (0, 0, tile_width, tile_height), + subpixels=subpixels, + fill_rule=entries[index][3], + ) + if analytic is not None: + return restore_tile(analytic[0], left, top) alpha = _fill_path_coverage( [control - offset for control in path], (0, 0, tile_width, tile_height), @@ -1763,6 +1898,24 @@ def rasterise_multi_group( # seed. The two-pixel antialias margin already makes these fixed tiles # conservative for the local coordinate updates used by the fit. initial_simple_groups = cropped_simple_groups() + # Simple paths use cropped tiles, so unlike the full-canvas compositor + # their bounds are optimisation state. A SAMVG coordinate update can move + # a boundary outside its initial two-pixel antialias fringe; continuing to + # rasterise the old crop silently clips that fill and creates the holes and + # spikes visible in long fits. Keep one packed reference so the movement + # check is a single device reduction; rebuild the inexpensive Python tile + # grouping only after a meaningful move. + simple_tile_reference = control_storage.detach().clone() + + def refresh_simple_tiles() -> None: + nonlocal initial_simple_groups, simple_tile_reference + + movement = (control_storage.detach() - simple_tile_reference).abs().amax() + if float(movement) <= 1.0: + return + initial_simple_groups = cropped_simple_groups() + simple_tile_reference = control_storage.detach().clone() + initial_multi_tiles = { index: tile_for(path) for index, path in enumerate(controls) @@ -1849,6 +2002,8 @@ def rasterise_multi_group( if alphas[index] is None: alphas[index] = rasterise_multi(index, path) alpha_stack = torch.stack([alpha for alpha in alphas if alpha is not None]) + if alpha_values is not None: + alpha_stack = alpha_stack * alpha_values.clamp(0, 1)[:, None, None] composite = ( _compiled_opaque_fill_composite() if goal.is_cuda @@ -1860,14 +2015,152 @@ def rasterise_multi_group( else _composite_opaque_fills(alpha_stack, color_storage, under) ) loss = ((rendered - goal) ** 2).mean() - loss = ( - loss - + xing_weight - * (_xing_penalties(all_controls) * xing_contour_weights).sum() - ) + loss = loss + xing_weight * _xing_loss(all_controls) loss.backward() point_optimizer.step() colour_optimizer.step() + close_contours() + refresh_simple_tiles() + continue + + if sparse_replay: + # Dense replay previously saved a full alpha, pre-layer canvas and + # downstream-transparency map for every SVG path. Painter-order + # compositing is local to a path's coverage tile, so retain only + # those slices while keeping the current canvas/transparency as + # full images. This is algebraically the same replay derivative. + with torch.no_grad(): + coverages: list[tuple[Any, int, int] | None] = [None] * len(entries) + for ( + _shape, + fill_rule, + tile_width, + tile_height, + ), items in simple_groups.items(): + for index, alpha, left, top in rasterise_simple_tiles( + fill_rule, tile_width, tile_height, items + ): + coverages[index] = (alpha, left, top) + for index, path in enumerate(controls): + if coverages[index] is None: + coverages[index] = (rasterise_multi(index, path), 0, 0) + + opacity_values = ( + alpha_values.detach().clamp(0, 1) + if alpha_values is not None + else None + ) + stored_alphas: list[Any] = [] + before_tiles: list[Any] = [] + rendered = torch.zeros_like(goal) if under is None else under.clone() + for index, item in enumerate(coverages): + assert item is not None + alpha, left, top = item + if opacity_values is not None: + alpha = alpha * opacity_values[index] + bottom, right = top + alpha.shape[0], left + alpha.shape[1] + canvas = rendered[top:bottom, left:right] + before_tiles.append(canvas.clone()) + rendered[top:bottom, left:right] = ( + canvas * (1 - alpha[..., None]) + + color_storage[index].detach().clamp(0, 1) * alpha[..., None] + ) + stored_alphas.append(alpha) + + suffix_tiles: list[Any] = [None] * len(entries) + transparency = torch.ones( + (work_height, work_width), dtype=goal.dtype, device=device + ) + for index in range(len(entries) - 1, -1, -1): + item = coverages[index] + assert item is not None + alpha, left, top = item + bottom, right = top + alpha.shape[0], left + alpha.shape[1] + suffix = transparency[top:bottom, left:right] + suffix_tiles[index] = suffix.clone() + suffix.mul_(1 - stored_alphas[index]) + image_gradient = 2 * (rendered - goal) / rendered.numel() + + def sparse_layer_loss( + index: int, + alpha: Any, + left: int, + top: int, + *, + saved_coverages: list[tuple[Any, int, int] | None] = coverages, + saved_alphas: list[Any] = stored_alphas, + saved_suffixes: list[Any] = suffix_tiles, + saved_canvases: list[Any] = before_tiles, + gradient: Any = image_gradient, + ) -> Any: + item = saved_coverages[index] + assert item is not None + coverage, _stored_left, _stored_top = item + stored_alpha = saved_alphas[index] + suffix = saved_suffixes[index] + canvas = saved_canvases[index] + bottom, right = top + alpha.shape[0], left + alpha.shape[1] + gradient = gradient[top:bottom, left:right] + colour = color_storage[index] + colour_delta = colour.detach().clamp(0, 1) - canvas + alpha_gradient = (gradient * suffix[..., None] * colour_delta).sum( + dim=-1 + ) + opacity = ( + alpha_values[index].clamp(0, 1) + if alpha_values is not None + else None + ) + colour_gradient = ( + gradient * suffix[..., None] * stored_alpha[..., None] + ).sum(dim=(0, 1)) + geometry_loss = (alpha * alpha_gradient.detach()).sum() + if opacity is not None: + geometry_loss = geometry_loss * opacity + geometry_loss = ( + geometry_loss + + opacity * (coverage * alpha_gradient.detach()).sum() + ) + return ( + geometry_loss + + (colour.clamp(0, 1) * colour_gradient.detach()).sum() + ) + + for ( + _shape, + fill_rule, + tile_width, + tile_height, + ), items in simple_groups.items(): + # Sparse replay keeps only a tile-local graph, so it can + # batch more equal-size paths than the legacy dense replay. + # This reduces native coverage launches without increasing the + # full-canvas memory footprint. + batch_size = sparse_backward_batch_size(tile_width, tile_height) + for offset in range(0, len(items), batch_size): + loss = torch.zeros((), device=device) + for index, alpha, left, top in rasterise_simple_tiles( + fill_rule, + tile_width, + tile_height, + items[offset : offset + batch_size], + ): + loss = loss + sparse_layer_loss(index, alpha, left, top) + loss.backward() + simple_indices = { + index + for group in simple_groups.values() + for index, _left, _top in group + } + for index, path in enumerate(controls): + if index not in simple_indices: + alpha = rasterise_multi(index, path) + sparse_layer_loss(index, alpha, 0, 0).backward() + (xing_weight * _xing_loss(all_controls)).backward() + point_optimizer.step() + colour_optimizer.step() + close_contours() + refresh_simple_tiles() continue # First composite the exact same soft fills without recording an @@ -1900,8 +2193,15 @@ def rasterise_multi_group( before: list[Any] = [] rendered = torch.zeros_like(goal) if under is None else under + opacity_values = ( + alpha_values.detach().clamp(0, 1) if alpha_values is not None else None + ) + initial_coverages: list[Any | None] = initial_alphas.copy() for index, alpha in enumerate(initial_alphas): assert alpha is not None + if opacity_values is not None: + alpha = alpha * opacity_values[index] + initial_alphas[index] = alpha colour = color_storage[index] before.append(rendered) rendered = ( @@ -1926,6 +2226,7 @@ def layer_loss( alphas: list[Any | None] = initial_alphas, suffixes: list[Any | None] = downstream, canvases: list[Any] = before, + coverages: list[Any | None] = initial_coverages, gradient: Any = image_gradient, ) -> Any: stored_alpha = alphas[index] @@ -1935,12 +2236,20 @@ def layer_loss( colour = color_storage[index] colour_delta = colour.detach().clamp(0, 1) - canvases[index] alpha_gradient = (gradient * suffix[..., None] * colour_delta).sum(dim=-1) + opacity = ( + alpha_values[index].clamp(0, 1) if alpha_values is not None else None + ) colour_gradient = ( gradient * suffix[..., None] * stored_alpha[..., None] ).sum(dim=(0, 1)) - return (alpha * alpha_gradient.detach()).sum() + ( - colour.clamp(0, 1) * colour_gradient.detach() - ).sum() + geometry_loss = (alpha * alpha_gradient.detach()).sum() + if opacity is not None: + geometry_loss = geometry_loss * opacity + coverage = coverages[index] + assert coverage is not None + opacity_gradient = (coverage * alpha_gradient.detach()).sum() + geometry_loss = geometry_loss + opacity * opacity_gradient + return geometry_loss + (colour.clamp(0, 1) * colour_gradient.detach()).sum() # Backpropagate a bounded batch at a time. The compositing derivative # above accounts for all later opaque layers, so this has the same MSE @@ -1982,14 +2291,14 @@ def layer_loss( index, rasterise_multi(index, path), ).backward() - ( - xing_weight * (_xing_penalties(all_controls) * xing_contour_weights).sum() - ).backward() + (xing_weight * _xing_loss(all_controls)).backward() point_optimizer.step() colour_optimizer.step() + close_contours() + refresh_simple_tiles() coordinate_scale_cpu = coordinate_scale.cpu() - for index, ((element, _contours, _colour, _fill_rule), path) in enumerate( + for index, ((element, _contours, _colour, _fill_rule, _opacity), path) in enumerate( zip(entries, controls, strict=True) ): colour = color_storage[index] @@ -2002,6 +2311,11 @@ def layer_loss( round(float(v) * 255) for v in colour.detach().clamp(0, 1).cpu() ) element.set("fill", f"#{red:02x}{green:02x}{blue:02x}") + if alpha_values is not None: + element.set( + "fill-opacity", + f"{float(alpha_values[index].detach().clamp(0, 1).cpu()):.8g}", + ) return ET.tostring(root, encoding="unicode") @@ -2187,8 +2501,9 @@ def fit_opaque_fills_locally( selected_indices: set[int] | None = None, optimisation_long_side: int | None = 64, gpu_gate: Any = None, + learn_alpha: bool = False, ) -> str: - """Fit one spatially bounded opaque-fill group as a local-search move. + """Fit one spatially bounded fill group as a local-search move. Unlike the legacy stroke fitter this operates on complete filled shapes, including compound paths and holes. It deliberately keeps the 64px @@ -2233,6 +2548,7 @@ def fit_opaque_fills_locally( steps=steps, optimisation_long_side=optimisation_long_side, backdrop=backdrop, + learn_alpha=learn_alpha, ) fitted_root = ET.fromstring(fitted) fitted_by_index = dict(enumerate(fitted_root.iter())) @@ -2242,6 +2558,8 @@ def fit_opaque_fills_locally( updated = fitted_by_index[index] element.set("d", updated.get("d", "")) element.set("fill", updated.get("fill", element.get("fill", ""))) + if learn_alpha: + element.set("fill-opacity", updated.get("fill-opacity", "1")) return ET.tostring(original, encoding="unicode") @@ -2254,10 +2572,16 @@ def fit_filled_svg_bounded( maximum_paths: int = 16, gpu_gate: Any = None, measurements: list[dict[str, int | float]] | None = None, + learn_alpha: bool = False, + global_replay: bool = True, ) -> str: """Run one full SAMVG fill phase as bounded spatial coordinate descent. - ``steps`` is the per-group phase budget. Coordinate descent needs to give + ``steps`` is the per-group phase budget. ``global_replay`` uses the + sparse painter-order replay to give every path the dissertation's one + simultaneous Adam update per iteration without materialising a full alpha + stack. The older coordinate-descent path remains available for local + experiments. Coordinate descent needs to give every group the same fitting opportunity that it would have had in the original global graph; splitting that budget between groups loses detail. It consequently trades wall time for a strictly bounded differentiable @@ -2266,6 +2590,41 @@ def fit_filled_svg_bounded( """ if steps < 1: raise ValueError("steps must be positive") + if global_replay: + import xml.etree.ElementTree as ET + + started = perf_counter() + peak_before = 0 + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + peak_before = int(torch.cuda.max_memory_allocated()) + except ImportError: + torch = None # type: ignore[assignment] + with gpu_slot(gpu_gate): + fitted = fit_filled_svg( + svg, + target, + steps=steps, + learn_alpha=learn_alpha, + sparse_replay=True, + ) + if measurements is not None: + peak = peak_before + if torch is not None and torch.cuda.is_available(): + torch.cuda.synchronize() + peak = int(torch.cuda.max_memory_allocated()) + measurements.append( + { + "group": 0, + "paths": len(_fittable_fill_elements(ET.fromstring(svg))), + "seconds": perf_counter() - started, + "peak_cuda_bytes": peak, + } + ) + return fitted groups = fill_groups(svg, maximum_paths=maximum_paths) if not groups: raise UnsupportedPathError("no opaque filled cubic paths to optimise") @@ -2292,6 +2651,7 @@ def fit_filled_svg_bounded( selected_indices=group, optimisation_long_side=None, gpu_gate=gpu_gate, + learn_alpha=learn_alpha, ) if measurements is not None: peak = peak_before diff --git a/src/vectrify/refine/samvg.md b/src/vectrify/refine/samvg.md new file mode 100644 index 0000000..e8a98b3 --- /dev/null +++ b/src/vectrify/refine/samvg.md @@ -0,0 +1,201 @@ +# SAMVG algorithm reference + +This is a pseudocode reference for the two-phase method described in Chapter 3 +of Yiding Zhu's *SAMVG* dissertation. It is a behavioural specification for +Vectrify's SAMVG-inspired path, not a copy of unreleased research code. + +The ordering below matters. In particular, impact is scored for a complete +cleaned SAM mask before that mask is split into connected components for path +tracing. Coverage prompts and residual prompts solve different problems and +must not be conflated. + +## Parameters defined by the dissertation + +```text +AUTOMATIC_POINT_GRID = 32 x 32 +RESIDUAL_THRESHOLD = 0.784 +FIT_STEPS_PER_PHASE = 500 +``` + +The dissertation leaves the SAM checkpoint, SAM confidence/stability gates, +small-region and hole thresholds, impact threshold, circular-kernel radius, +and optimiser hyperparameters as implementation choices. Keep those as +explicit parameters and benchmark them; do not infer a canonical value from a +path-count target alone. + +## Reported representation variations + +The baseline uses a fixed number of cubic segments per closed contour and +opaque fills. The dissertation also reports two independent variations: + +```text +SAMVG+var = select locally distinct contour points whose curvature score + crosses a caller-selected threshold, then fit one cubic between + each adjacent selected pair +SAMVG+alpha = make every path fill opacity an optimisation parameter +``` + +These are representation changes, not mask-selection changes. They must be +enabled explicitly when comparing with an SVG made by either variation; the +threshold value itself is not specified by the dissertation. + +## Data types + +```text +Mask = boolean H x W image +PaintedMask = (mask: Mask, colour: RGB) +Path = closed filled SVG path +Document = ordered list of SVG elements +Canvas = RGB H x W image +``` + +`Composite(canvas, mask, colour)` paints `colour` wherever `mask` is true. +`Error(target, canvas)` is the pixel reconstruction error used consistently +within a filtering pass. For a blank initial canvas, uncovered pixels receive +the maximum error so the first mask is not biased toward bright colours. + +## Common mask preparation and impact filter + +```text +function CLEAN_MASKS(raw_masks): + # This is SAM automatic-mask post-processing, before SAMVG selection. + masks = retain masks passing SAM's predicted-quality/stability gates + masks = remove configured small connected regions and small holes + return masks + + +function FILTER_BY_IMPACT(target, masks, initial_canvas, min_improvement): + # Each entry remains a WHOLE cleaned SAM mask until it is accepted. + candidates = CLEAN_MASKS(masks) + candidates = sort candidates by descending mask area + + canvas = copy(initial_canvas) + accepted = [] + + for mask in candidates: + colour = mean_rgb(target pixels where mask is true) + proposal = Composite(canvas, mask, colour) + + improvement = Error(target, canvas) - Error(target, proposal) + if improvement < min_improvement: + continue + + accepted.append((mask, colour, improvement)) + canvas = proposal + + return accepted, canvas +``` + +Do **not** score disconnected components of one SAM mask independently. Split +an accepted mask only when tracing it: every connected component becomes its +own editable SVG path, inherits the accepted mask colour, and preserves the +accepted mask's painter-order slot. + +```text +function TRACE_ACCEPTED_MASKS(accepted_masks): + document = [] + for (mask, colour, _) in accepted_masks in acceptance order: + for component in connected_components(mask): + contour = extract_outer_contour(component) + path = fit_fixed_segment_bezier_path(contour) + document.append(filled_path(path, colour)) + return document +``` + +## Phase 1: segmentation, coverage recovery, and first fit + +```text +function FIRST_PHASE(target): + raw = SAM_AUTOMATIC_MASK_GENERATION( + target, + point_grid = AUTOMATIC_POINT_GRID, + crop_schedule = SAM_AMG_CROP_SCHEDULE, + ) + # AMG removes duplicate candidates in two passes: predicted-IoU NMS + # within each crop, then crop-area-priority NMS across all crop outputs. + # The latter prefers a duplicate from the smaller crop. + + # Begin on a blank canvas and retain useful whole masks in area order. + first_masks, mask_canvas = FILTER_BY_IMPACT( + target, raw, blank_canvas(target.size), min_improvement + ) + + # This recovery looks for regions with no retained-mask coverage. It is + # still part of segmentation, before any SVG path optimisation. + uncovered = NOT union(mask for (mask, _, _) in first_masks) + coverage_map = circular_convolution(uncovered) + coverage_candidates = coordinates_of_full_empty_circles(coverage_map) + coverage_centres = mean_shift_clusters(coverage_candidates) + coverage_raw = SAM_PROMPTED_MASKS(target, coverage_centres) + + # Score newly prompted masks against the retained-mask composite, not a + # fresh blank canvas, and append accepted masks after existing masks. + coverage_masks, _ = FILTER_BY_IMPACT( + target, coverage_raw, mask_canvas, min_improvement + ) + seed = TRACE_ACCEPTED_MASKS(first_masks + coverage_masks) + + first_fit = OPTIMISE_PATHS(seed, target, steps = FIT_STEPS_PER_PHASE) + return seed, first_fit +``` + +The coverage pass finds *uncovered* areas. It cannot recover texture inside a +large filled mask that already covers the relevant pixels. + +## Phase 2: residual-detail recovery and second fit + +```text +function SECOND_PHASE(target, first_fit): + fitted_canvas = RASTERISE(first_fit) + difference = sum_over_rgb(abs(target - fitted_canvas)) + + # Unlike coverage recovery, this detects high-error regions after fitting, + # including regions that are already alpha-covered by a coarse fill. + residual_map = circular_convolution(difference) + residual_regions = connected_components( + threshold(residual_map, RESIDUAL_THRESHOLD) + ) + residual_centres = centres(residual_regions) + residual_raw = SAM_PROMPTED_MASKS(target, residual_centres) + + # Filter against the fitted render so only new masks that reduce remaining + # error are kept. Append their paths in painter order after first_fit. + residual_masks, _ = FILTER_BY_IMPACT( + target, residual_raw, fitted_canvas, min_improvement + ) + additions = TRACE_ACCEPTED_MASKS(residual_masks) + + recovered_document = append_in_painter_order(first_fit, additions) + final_fit = OPTIMISE_PATHS( + recovered_document, target, steps = FIT_STEPS_PER_PHASE + ) + return recovered_document, final_fit +``` + +## End-to-end procedure + +```text +function SAMVG(target): + seed, first_fit = FIRST_PHASE(target) + recovered_document, final_fit = SECOND_PHASE(target, first_fit) + return { + seed, + first_fit, + recovered_document, + final_fit, + } +``` + +## Implementation invariants + +- Preserve document/painter order throughout; accepted additions come after + the canvas against which their impact was measured. +- Use the same cleanup, complete-mask impact filtering, and fixed-segment + tracing procedure for automatic, coverage-prompted, and residual-prompted + masks. +- Keep coverage recovery and residual recovery as separate passes. A high + coverage percentage does not show that residual detail recovery is needless. +- Judge a run by raster error and visual output, not just retained-mask or + emitted-path count. Component splitting can make these counts diverge. +- Optional stroke/text handling is an extension around this fill pipeline; it + must not alter the fill-mask ordering or residual-recovery criteria. diff --git a/src/vectrify/refine/samvg.py b/src/vectrify/refine/samvg.py index fa0b8b3..7e22b4c 100644 --- a/src/vectrify/refine/samvg.py +++ b/src/vectrify/refine/samvg.py @@ -39,11 +39,20 @@ # grid. 64 doubles the old 32 while leaving full-resolution-mask # headroom on a 16 GB GPU; users with larger cards can raise it by environment. SAMVG_POINTS_PER_BATCH = int(os.environ.get("VECTRIFY_SAMVG_POINTS_PER_BATCH", "64")) -# SAMVG's own impact filter selects useful masks against the image. Retaining -# AMG's score gates here discarded the small facial candidates needed by the -# photo seed before that image-aware test could evaluate them. -SAMVG_PRED_IOU_THRESH = 0.0 -SAMVG_STABILITY_SCORE_THRESH = 0.0 +# SAMVG's image-aware impact filter is the retained-mask decision specified by +# the dissertation. Keep AMG's confidence gates configurable, but disable +# them by default so a small, useful candidate reaches that later test instead +# of being discarded by a checkpoint-confidence heuristic. +SAMVG_PRED_IOU_THRESH = float(os.environ.get("VECTRIFY_SAMVG_PRED_IOU_THRESH", "0")) +SAMVG_STABILITY_SCORE_THRESH = float( + os.environ.get("VECTRIFY_SAMVG_STABILITY_SCORE_THRESH", "0") +) +# The dissertation specifies a fixed circular residual kernel scaled to the +# image, but not its fraction. Cat calibration selects this value by final +# raster error and complexity; callers can reproduce alternate sweeps. +SAMVG_RESIDUAL_RADIUS_FRACTION = float( + os.environ.get("VECTRIFY_SAMVG_RESIDUAL_RADIUS_FRACTION", "0.005") +) # The SAMVG seed only needs OCR once and does it after SAM has released its # automatic-mask pipeline. This is a real VLM pass, not a separate small OCR # detector: it can decide which visible labels deserve editable text and place @@ -460,7 +469,7 @@ class _SamRuntime: embedding_size: tuple[int, int] | None = None -def _sam_runtime() -> _SamRuntime: +def _sam_runtime(*, model: str = SAMVG_MODEL) -> _SamRuntime: """Load SAM once, in half precision when CUDA is available.""" try: import torch @@ -469,13 +478,13 @@ def _sam_runtime() -> _SamRuntime: raise ImportError( "SAMVG requires the samvg extra. Install 'vectrify[samvg]'." ) from exc - options: dict[str, Any] = {"model": SAMVG_MODEL, "device": 0} + options: dict[str, Any] = {"model": model, "device": 0} if torch.cuda.is_available(): options["dtype"] = torch.float16 generator = pipeline("mask-generation", **options) log.info( "SAMVG automatic masks: %s on %s (%s).", - SAMVG_MODEL, + model, generator.device, "fp16" if torch.cuda.is_available() else "fp32", ) @@ -492,61 +501,122 @@ def _sam_autocast(): def _automatic_forward(inputs: Any, runtime: _SamRuntime) -> dict[str, Any]: - """Decode on CUDA, then expand and filter masks on CPU. + """Decode and filter one AMG prompt batch in SAM's original order.""" + import torch - The stock Transformers pipeline expands a prompt batch to the original - image size on CUDA. At 1024px that transient allocation is larger than the - decoder itself. Its filtering sequence is unchanged here; only the - post-decoder device changes. - """ generator = runtime.generator - input_boxes = inputs.pop("input_boxes").detach().cpu().float() + input_boxes = inputs.pop("input_boxes").float() is_last = inputs.pop("is_last") original_sizes = inputs.pop("original_sizes").detach().cpu().tolist() reshaped_sizes = inputs.pop("reshaped_input_sizes", None) if reshaped_sizes is not None: reshaped_sizes = reshaped_sizes.detach().cpu().tolist() - with _sam_autocast(): + # `.cpu()` alone preserves the decoder's autograd graph, retaining every + # prior prompt batch's CUDA activations. AMG is inference-only, so make + # that lifetime explicit before handing compact candidates to the host. + with torch.inference_mode(), _sam_autocast(): model_outputs = generator.model(**inputs) - masks = generator.image_processor.post_process_masks( - model_outputs.pred_masks.detach().cpu(), - original_sizes, - mask_threshold=0, - reshaped_input_sizes=reshaped_sizes, - binarize=False, - ) - filtered_masks, scores, boxes = generator.image_processor.filter_masks( - masks[0], - model_outputs.iou_scores.detach().cpu().float()[0], - original_sizes[0], - input_boxes[0], - SAMVG_PRED_IOU_THRESH, - SAMVG_STABILITY_SCORE_THRESH, - 0, - 1, - ) + # Official AMG interpolates decoder logits to the crop canvas before + # its confidence, stability, and crop-edge tests. Filtering at 256px + # is faster but changes which fine masks survive, so it cannot be used + # for the dissertation-faithful seed. + masks_at_size = generator.image_processor.post_process_masks( + model_outputs.pred_masks, + original_sizes, + reshaped_input_sizes=reshaped_sizes, + mask_threshold=0, + binarize=False, + )[0] + masks, scores, boxes = _filter_automatic_masks( + masks_at_size, + model_outputs.iou_scores[0], + original_sizes[0], + input_boxes[0], + ) return { - "masks": filtered_masks, + "masks": masks, "is_last": is_last, "boxes": boxes, - "iou_scores": scores, + "scores": scores, + "original_size": original_sizes[0], + "reshaped_size": reshaped_sizes[0] if reshaped_sizes is not None else None, + "crop_box": input_boxes[0], } -def _automatic_masks_for( +def _filter_automatic_masks( + masks: Any, + iou_scores: Any, + original_size: list[int], + cropped_box_image: Any, +) -> tuple[Any, Any, Any]: + """Apply SAM AMG's full-resolution confidence and crop-edge filtering.""" + import torch + from transformers.models.sam.image_processing_sam import ( + _batched_mask_to_box, + _compute_stability_score, + _is_box_near_crop_edge, + _pad_masks, + ) + + original_height, original_width = original_size + scores = iou_scores.reshape(-1) + masks = masks.reshape(-1, *masks.shape[-2:]) + keep = torch.ones(len(masks), dtype=torch.bool, device=masks.device) + if SAMVG_PRED_IOU_THRESH > 0: + keep &= scores > SAMVG_PRED_IOU_THRESH + if SAMVG_STABILITY_SCORE_THRESH > 0: + keep &= _compute_stability_score(masks, 0, 1) > SAMVG_STABILITY_SCORE_THRESH + masks, scores = masks[keep] > 0, scores[keep] + boxes = _batched_mask_to_box(masks) + keep = ~_is_box_near_crop_edge( + boxes, cropped_box_image, [0, 0, original_width, original_height] + ) + return ( + _pad_masks(masks[keep], cropped_box_image, original_height, original_width), + scores[keep], + boxes[keep], + ) + + +def _nms_indices(boxes: Any, scores: Any) -> Any: + """Use SAM AMG's box-NMS configuration with dtype-safe scores.""" + import torch + from torchvision.ops import batched_nms + + return batched_nms( + boxes=boxes.float(), + scores=scores.float(), + idxs=torch.zeros(len(boxes), dtype=torch.long), + iou_threshold=0.7, + ) + + +def _finalize_automatic_masks(masks: Any, scores: Any, boxes: Any) -> list[np.ndarray]: + """Apply AMG's final image-global NMS and transfer binary masks.""" + if not len(masks): + return [] + keep = _nms_indices(boxes, scores) + selected_masks = masks[keep] + return [np.asarray(mask.cpu(), dtype=bool) for mask in selected_masks] + + +def _automatic_mask_candidates_for( source: Image.Image, runtime: _SamRuntime, *, cache_embedding: bool, points_per_batch: int = SAMVG_POINTS_PER_BATCH, -) -> list[np.ndarray]: - """Run one AMG image/crop without recomputing prompt-grid embeddings. +) -> tuple[Any, Any, Any]: + """Return post-filter AMG candidates before its image-global crop NMS. Transformers' public mask-generation call already encodes an image once per 32x32 prompt grid. For the full image we use the same pipeline stages directly so the resulting embedding can be reused by coverage/residual prompts. Crops intentionally retain their own embeddings. """ + import torch + generator = runtime.generator arguments = { "points_per_batch": points_per_batch, @@ -558,7 +628,22 @@ def _automatic_masks_for( # Keep a small compatibility path for mocked/older Transformers pipelines. if not hasattr(generator, "preprocess"): output = generator(source, **arguments) - return [np.asarray(mask, dtype=bool) for mask in output["masks"]] + masks = ( + torch.from_numpy( + np.stack([np.asarray(mask, dtype=bool) for mask in output["masks"]]) + ) + if output["masks"] + else torch.empty((0, source.height, source.width), dtype=torch.bool) + ) + boxes = torch.empty((len(masks), 4), dtype=torch.float32) + for index, mask in enumerate(masks): + ys, xs = torch.where(mask) + boxes[index] = torch.tensor( + (xs.min(), ys.min(), xs.max() + 1, ys.max() + 1), dtype=torch.float32 + ) + scores = torch.ones(len(masks)) + keep = _nms_indices(boxes, scores) if len(masks) else [] + return masks[keep], scores[keep], boxes[keep] outputs = [] for inputs in generator.preprocess( @@ -579,14 +664,34 @@ def _automatic_masks_for( runtime.image_embeddings = embedding runtime.embedding_size = source.size outputs.append(_automatic_forward(inputs, runtime)) - output = generator.postprocess(outputs) - return [np.asarray(mask, dtype=bool) for mask in output["masks"]] + # Keep the GPU bounded to one decoder batch. NMS and score filtering + # have already retained only the decoder logits that need the + # full-resolution SAM post-processing. + outputs[-1]["masks"] = outputs[-1]["masks"].cpu() + outputs[-1]["scores"] = outputs[-1]["scores"].cpu() + outputs[-1]["boxes"] = outputs[-1]["boxes"].cpu() + masks = [output["masks"] for output in outputs if len(output["masks"])] + if not masks: + return ( + torch.empty((0, source.height, source.width), dtype=torch.bool), + torch.empty(0), + torch.empty((0, 4)), + ) + masks = torch.cat(masks) + scores = torch.cat([output["scores"] for output in outputs if len(output["masks"])]) + boxes = torch.cat([output["boxes"] for output in outputs if len(output["masks"])]) + # Meta's AMG first suppresses candidates within each crop by predicted + # quality. Its second, cross-crop pass happens in ``automatic_masks`` + # below and deliberately ranks the surviving masks by crop area instead. + keep = _nms_indices(boxes, scores) + return masks[keep], scores[keep], boxes[keep] def automatic_masks( image: Image.Image, *, max_side: int | None = SAMVG_MAX_SIDE, + points_per_batch: int = SAMVG_POINTS_PER_BATCH, _runtime: _SamRuntime | None = None, ) -> list[np.ndarray]: """Retrieve SAM AMG masks with the thesis grid, optionally size-capped.""" @@ -601,37 +706,54 @@ def automatic_masks( width, height = image.size def collect(points_per_batch: int) -> list[np.ndarray]: - collected = _automatic_masks_for( + import torch + + masks, scores, boxes = _automatic_mask_candidates_for( image, runtime, cache_embedding=True, points_per_batch=points_per_batch, ) + all_masks = [masks] + all_scores = [torch.full_like(scores, 1 / (width * height))] + all_boxes = [boxes] overlap = int((512 / 1500) * min(width, height)) crop_width = math.ceil((overlap + width) / 2) crop_height = math.ceil((overlap + height) / 2) - for x, y in { + for x, y in ( (0, 0), - (crop_width - overlap, 0), (0, crop_height - overlap), + (crop_width - overlap, 0), (crop_width - overlap, crop_height - overlap), - }: + ): right, bottom = min(x + crop_width, width), min(y + crop_height, height) crop_box = (x, y, right, bottom) - for crop_mask in _automatic_masks_for( + crop_masks, crop_scores, crop_boxes = _automatic_mask_candidates_for( image.crop(crop_box), runtime, cache_embedding=False, points_per_batch=points_per_batch, - ): - if _is_crop_edge_mask(crop_mask, crop_box, image.size): + ) + for index, crop_mask in enumerate(crop_masks): + crop_mask_array = np.asarray(crop_mask, dtype=bool) + if _is_crop_edge_mask(crop_mask_array, crop_box, image.size): continue mask = np.zeros((height, width), dtype=bool) - mask[y:bottom, x:right] = crop_mask - collected.append(mask) - return collected + mask[y:bottom, x:right] = crop_mask_array + all_masks.append(torch.from_numpy(mask)[None]) + crop_area = (right - x) * (bottom - y) + all_scores.append( + torch.full_like(crop_scores[index : index + 1], 1 / crop_area) + ) + box = crop_boxes[index].clone() + box[[0, 2]] += x + box[[1, 3]] += y + all_boxes.append(box[None]) + return _finalize_automatic_masks( + torch.cat(all_masks), torch.cat(all_scores), torch.cat(all_boxes) + ) - collected = collect(SAMVG_POINTS_PER_BATCH) + collected = collect(points_per_batch) return [_restore_mask(mask, original_size) for mask in collected] @@ -652,25 +774,44 @@ def _components( for runs in _run_components(foreground): if sum(end - start for _y, start, end in runs) < min_pixels: continue + min_y = min(y for y, _start, _end in runs) + max_y = max(y for y, _start, _end in runs) + min_x = min(start for _y, start, _end in runs) + max_x = max(end for _y, _start, end in runs) + local = np.zeros((max_y - min_y + 1, max_x - min_x), dtype=bool) + for y, start, end in runs: + local[y - min_y, start - min_x : end - min_x] = True component = np.zeros((height, width), dtype=bool) for y, start, end in runs: component[y, start:end] = True - if fill_holes: + # A hole must contain at least one non-component pixel strictly inside + # this box. Most small SAM fragments are solid or only touch the box + # boundary, so avoid a connected-components pass when a hole is + # impossible. + has_interior_background = ( + local.shape[0] > 2 and local.shape[1] > 2 and not local[1:-1, 1:-1].all() + ) + if fill_holes and has_interior_background: # AMG's postprocessing removes *small* enclosed holes, rather # than turning meaningful cutouts such as an eye into a solid # region. The same area cutoff as tiny components keeps those # two decisions consistent. - for hole in _run_components(~component): + # The exterior background necessarily reaches a component bounding + # box edge, while an enclosed hole cannot. Checking this compact + # box is equivalent to checking the full mask, without scanning a + # 1024px canvas once for every small disconnected component. + local_height, local_width = local.shape + for hole in _run_components(~local): area = sum(end - start for _y, start, end in hole) if area > min_pixels: continue touches_border = any( - y in {0, height - 1} or start == 0 or end == width + y in {0, local_height - 1} or start == 0 or end == local_width for y, start, end in hole ) if not touches_border: for y, start, end in hole: - component[y, start:end] = True + component[y + min_y, start + min_x : end + min_x] = True components.append(np.asarray(component, dtype=bool)) return components @@ -740,7 +881,7 @@ def filter_by_impact( initial_canvas: np.ndarray | None = None, initial_coverage: np.ndarray | None = None, min_pixels: int = 32, - min_impact: float = 1e-5, + min_impact: float = 3e-6, max_layers: int = 128, fill_holes: bool = True, ) -> list[MaskLayer]: @@ -765,19 +906,24 @@ def filter_by_impact( error_map = _impact_error_map(target, canvas, coverage) error_total = float(error_map.sum(dtype=np.float64)) error = error_total / error_map.size - initial_count = len(accepted) - candidates = [ - component - for mask in masks - if np.asarray(mask).shape == (height, width) - for component in _components( - np.asarray(mask, dtype=bool), min_pixels, fill_holes=fill_holes + # SAMVG filters an AMG *mask* by its rendered impact, after AMG's component + # cleanup. Components are independent paths only in the subsequent tracing + # stage. Scoring every disconnected component here changes the paper's + # painter-order decision and promotes low-information rectangular fragments. + candidates = [] + for raw_mask in masks: + if np.asarray(raw_mask).shape != (height, width): + continue + components = _components( + np.asarray(raw_mask, dtype=bool), min_pixels, fill_holes=fill_holes ) - ] - candidates.sort(key=lambda mask: int(mask.sum()), reverse=True) - for mask in candidates: - if int(mask.sum()) < min_pixels: + if not components: continue + mask = np.logical_or.reduce(components) + candidates.append((mask, components)) + candidates.sort(key=lambda candidate: int(candidate[0].sum()), reverse=True) + retained: list[tuple[list[np.ndarray], tuple[int, int, int], float]] = [] + for mask, components in candidates: colour = cast( tuple[int, int, int], tuple(int(value) for value in np.rint(target[mask].mean(axis=0))), @@ -793,7 +939,7 @@ def filter_by_impact( impact = error - next_error if impact < min_impact: continue - accepted.append(MaskLayer(mask, colour, impact)) + retained.append((components, colour, impact)) canvas[mask] = colour coverage |= mask error_map[mask] = next_error_values @@ -801,8 +947,12 @@ def filter_by_impact( # Each SAMVG stage is allowed its own retained-mask budget. Applying # this to the combined existing+new list silently limited recovery to # one path once the automatic stage had filled its budget. - if len(accepted) - initial_count >= max_layers: + if len(retained) >= max_layers: break + for components, colour, impact in retained: + accepted.extend( + MaskLayer(component, colour, impact) for component in components + ) return accepted @@ -811,7 +961,7 @@ def coverage_prompt_points( shape: tuple[int, int], *, radius_fraction: float = 0.06, - max_points: int = 16, + max_points: int | None = None, ) -> list[tuple[int, int]]: """Find mean-shift centres of large circles untouched by retained masks.""" _canvas, coverage = _render_layers(shape, layers) @@ -827,7 +977,44 @@ def coverage_prompt_points( ((float(distance[round(y), round(x)]), round(x), round(y)) for x, y in centres), reverse=True, ) - return [(x, y) for _distance, x, y in ranked[:max_points]] + selected = ranked if max_points is None else ranked[:max_points] + return [(x, y) for _distance, x, y in selected] + + +def _circular_component_centres( + values: np.ndarray, + radius: int, + *, + threshold: float, + max_points: int | None = None, +) -> list[tuple[int, int]]: + """Return ranked centres of thresholded circular-convolution components.""" + import torch + import torch.nn.functional as functional + + if radius < 1: + raise ValueError("radius must be positive") + yy, xx = np.ogrid[-radius : radius + 1, -radius : radius + 1] + kernel = (xx * xx + yy * yy <= radius * radius).astype(np.float32) + padded = np.pad(np.asarray(values, dtype=np.float32), radius, mode="symmetric") + smoothed = functional.conv2d( + torch.from_numpy(padded)[None, None], + torch.from_numpy((kernel / kernel.sum())[None, None]), + )[0, 0].numpy() + labels, count = _label(smoothed >= threshold) + ranked: list[tuple[float, int, int]] = [] + for index in range(1, count + 1): + ys, xs = np.nonzero(labels == index) + if len(xs): + # The mean is the component centre prescribed by SAMVG. Ranking + # by response is deterministic when callers cap prompt count. + ranked.append( + (float(smoothed[ys, xs].mean()), round(xs.mean()), round(ys.mean())) + ) + selected = sorted(ranked, reverse=True) + if max_points is not None: + selected = selected[:max_points] + return [(x, y) for _score, x, y in selected] def prompted_masks( @@ -835,6 +1022,7 @@ def prompted_masks( points: list[tuple[int, int]], *, max_side: int | None = SAMVG_MAX_SIDE, + points_per_batch: int = SAMVG_POINTS_PER_BATCH, _runtime: _SamRuntime | None = None, ) -> list[np.ndarray]: """Prompt SAM at centres and return all three masks per point. @@ -859,32 +1047,36 @@ def prompted_masks( if runtime.processor is None: runtime.processor = SamProcessor(runtime.generator.image_processor) try: - input_points = [[[list(point)] for point in points]] - inputs = runtime.processor( - images=image, input_points=input_points, return_tensors="pt" - ).to(device) - if ( - runtime.embedding_size == image.size - and runtime.image_embeddings is not None - ): - # The full-image automatic pass has already encoded these pixels. - # Retain only decoder inputs for the coverage/residual prompts. - inputs.pop("pixel_values") - inputs["image_embeddings"] = runtime.image_embeddings - with torch.inference_mode(), _sam_autocast(): - output = runtime.generator.model(**inputs) - post = runtime.processor.image_processor.post_process_masks( - output.pred_masks.detach().cpu(), - inputs["original_sizes"].detach().cpu(), - inputs["reshaped_input_sizes"].detach().cpu(), - )[0] - return [ - _restore_mask( - np.asarray(post[prompt, candidate], dtype=bool), original_size + output_masks = [] + for start in range(0, len(points), points_per_batch): + batch = points[start : start + points_per_batch] + input_points = [[[list(point)] for point in batch]] + inputs = runtime.processor( + images=image, input_points=input_points, return_tensors="pt" + ).to(device) + if ( + runtime.embedding_size == image.size + and runtime.image_embeddings is not None + ): + # The full-image automatic pass has already encoded these pixels. + # Retain only decoder inputs for the coverage/residual prompts. + inputs.pop("pixel_values") + inputs["image_embeddings"] = runtime.image_embeddings + with torch.inference_mode(), _sam_autocast(): + output = runtime.generator.model(**inputs) + post = runtime.processor.image_processor.post_process_masks( + output.pred_masks.detach().cpu(), + inputs["original_sizes"].detach().cpu(), + inputs["reshaped_input_sizes"].detach().cpu(), + )[0] + output_masks.extend( + _restore_mask( + np.asarray(post[prompt, candidate], dtype=bool), original_size + ) + for prompt in range(post.shape[0]) + for candidate in range(post.shape[1]) ) - for prompt in range(post.shape[0]) - for candidate in range(post.shape[1]) - ] + return output_masks finally: if own_runtime and torch.cuda.is_available(): torch.cuda.empty_cache() @@ -895,18 +1087,25 @@ def retrieve_layers( masks: list[np.ndarray] | None = None, *, min_pixels: int = 32, - min_impact: float = 1e-5, + min_impact: float = 3e-6, max_layers: int = 512, fill_holes: bool = True, max_side: int | None = SAMVG_MAX_SIDE, + model: str = SAMVG_MODEL, + points_per_batch: int = SAMVG_POINTS_PER_BATCH, _runtime: _SamRuntime | None = None, ) -> list[MaskLayer]: """Run SAMVG's automatic-mask, coverage-prompt, filter sequence.""" image = image.convert("RGB") runtime = _runtime if masks is None: - runtime = runtime or _sam_runtime() - initial = automatic_masks(image, max_side=max_side, _runtime=runtime) + runtime = runtime or _sam_runtime(model=model) + initial = automatic_masks( + image, + max_side=max_side, + points_per_batch=points_per_batch, + _runtime=runtime, + ) else: initial = masks layers = filter_by_impact( @@ -917,9 +1116,14 @@ def retrieve_layers( max_layers=max_layers, fill_holes=fill_holes, ) - layers = recolour_visible_layers(image, layers) points = coverage_prompt_points(layers, (image.height, image.width)) - prompted = prompted_masks(image, points, max_side=max_side, _runtime=runtime) + prompted = prompted_masks( + image, + points, + max_side=max_side, + points_per_batch=points_per_batch, + _runtime=runtime, + ) recovered = filter_by_impact( image, prompted, @@ -929,7 +1133,6 @@ def retrieve_layers( max_layers=max_layers, fill_holes=fill_holes, ) - recovered = recolour_visible_layers(image, recovered) log.info( "SAMVG first pass: %d automatic mask(s), %d retained; %d coverage " "prompt(s), %d prompted mask(s), %d total retained.", @@ -939,7 +1142,12 @@ def retrieve_layers( len(prompted), len(recovered), ) - return recovered + # Mask selection intentionally scores the initially painted colours: that + # is the paper's greedy impact procedure. Once painter order is fixed, + # however, a lower layer should be coloured from only the pixels it still + # exposes. This is the least-squares fill for the emitted seed and does + # not alter its accepted masks, ordering, or coverage prompts. + return recolour_visible_layers(image, recovered) def _loops(mask: np.ndarray) -> list[list[tuple[float, float]]]: @@ -972,18 +1180,24 @@ def _loops(mask: np.ndarray) -> list[list[tuple[float, float]]]: return loops -def _corners(loop: list[tuple[float, float]], count: int) -> list[int]: - """Global curvature maxima with the local exclusion SAMVG describes.""" +def _curvature_scores(loop: list[tuple[float, float]]) -> np.ndarray: + """Return SAMVG's scale-aware cosine curvature score for a contour.""" points = np.asarray(loop, dtype=np.float32) size = len(points) - count = min(count, size) step = max(1, size // 12) before = points - np.roll(points, step, axis=0) after = np.roll(points, -step, axis=0) - points denom = np.linalg.norm(before, axis=1) * np.linalg.norm(after, axis=1) - score = np.divide( + return np.divide( (before * after).sum(axis=1), denom, out=np.ones(size), where=denom > 0 ) + + +def _corners(loop: list[tuple[float, float]], count: int) -> list[int]: + """Global curvature maxima with the local exclusion SAMVG describes.""" + size = len(loop) + count = min(count, size) + score = _curvature_scores(loop) blocked = np.zeros(size, dtype=bool) chosen: list[int] = [] exclusion = max(1, size // (count * 2)) @@ -1000,6 +1214,40 @@ def _corners(loop: list[tuple[float, float]], count: int) -> list[int]: return sorted(chosen) +def _variable_corners( + loop: list[tuple[float, float]], *, threshold: float, maximum: int +) -> list[int]: + """Select local curvature extrema below SAMVG+var's threshold. + + The dissertation's variable-segment variation replaces the fixed top-N + selection with a curvature threshold. Its threshold is not published, so + callers must choose it explicitly. ``maximum`` is only a safety bound for + pathological raster staircases, not a target complexity. + """ + size = len(loop) + if size < 3: + return [] + score = _curvature_scores(loop) + # SAMVG+var reverts the fixed variant's global-maxima-with-exclusion rule + # to the conventional local-extrema selector. The curvature *score* + # itself uses k-neighbours (Eq. 3-4); expanding the extrema neighbourhood + # to that same k suppresses genuine nearby corners and is not part of the + # variable-segment procedure. The asymmetric comparison retains one + # representative for a flat raster-corner plateau without coalescing + # separate extrema. + previous = np.roll(score, 1) + following = np.roll(score, -1) + local_minimum = (score < previous) & (score <= following) + eligible = np.flatnonzero(local_minimum & (score <= threshold)) + if len(eligible) < 3: + return _corners(loop, min(3, size)) + return ( + sorted(int(index) for index in eligible[:maximum]) + if len(eligible) >= 3 + else _corners(loop, min(3, size)) + ) + + def _fit_cubic( points: np.ndarray, *, reparameterize: bool = True ) -> tuple[np.ndarray, np.ndarray]: @@ -1072,11 +1320,23 @@ def solve(parameters: np.ndarray) -> np.ndarray: return controls[0], controls[1] -def _cubic_loop(loop: list[tuple[float, float]], segments: int) -> str | None: +def _cubic_loop( + loop: list[tuple[float, float]], + segments: int, + *, + curvature_threshold: float | None = None, + maximum_segments: int = 2048, +) -> str | None: size = len(loop) if size < 3: return None - corners = _corners(loop, segments) + corners = ( + _corners(loop, segments) + if curvature_threshold is None + else _variable_corners( + loop, threshold=curvature_threshold, maximum=maximum_segments + ) + ) if len(corners) < 3: return None points = np.asarray(loop, dtype=np.float32) @@ -1086,7 +1346,10 @@ def _cubic_loop(loop: list[tuple[float, float]], segments: int) -> str | None: np.arange(first, second + 1 if second >= first else second + size + 1) % size ) - sample = np.vstack((points[indices], points[second])) + # ``indices`` already includes the endpoint. Repeating it adds an + # artificial least-squares weight at every selected corner and bends + # each fitted cubic toward its end point rather than the contour data. + sample = points[indices] control_a, control_b = _fit_cubic(sample) end = points[second] parts.append( @@ -1097,12 +1360,28 @@ def _cubic_loop(loop: list[tuple[float, float]], segments: int) -> str | None: def mask_path( - mask: np.ndarray, *, segments: int = 8, overlap_pixels: int = 0 + mask: np.ndarray, + *, + segments: int = 8, + overlap_pixels: int = 0, + curvature_threshold: float | None = None, + maximum_segments: int = 2048, ) -> str | None: - """Fit every mask contour as a fixed-count cubic Bezier SVG path.""" + """Fit every mask contour as fixed-count or thresholded cubic Beziers.""" if overlap_pixels: mask = _binary_dilation(mask, overlap_pixels) - parts = [piece for loop in _loops(mask) if (piece := _cubic_loop(loop, segments))] + parts = [ + piece + for loop in _loops(mask) + if ( + piece := _cubic_loop( + loop, + segments, + curvature_threshold=curvature_threshold, + maximum_segments=maximum_segments, + ) + ) + ] return " ".join(parts) or None @@ -1352,7 +1631,12 @@ def mask_strokes( def _layer_svg_attributes( - layer: MaskLayer, segments: int, *, hybrid_strokes: bool = True + layer: MaskLayer, + segments: int, + *, + hybrid_strokes: bool = True, + curvature_threshold: float | None = None, + maximum_segments: int = 2048, ) -> list[dict[str, str]]: """Trace one SAM mask, using optional strokes only outside the thesis mode.""" colour = f"#{layer.colour[0]:02x}{layer.colour[1]:02x}{layer.colour[2]:02x}" @@ -1373,7 +1657,13 @@ def _layer_svg_attributes( } for data, width in strokes ] - data = mask_path(layer.mask, segments=segments, overlap_pixels=layer.overlap_pixels) + data = mask_path( + layer.mask, + segments=segments, + overlap_pixels=layer.overlap_pixels, + curvature_threshold=curvature_threshold, + maximum_segments=maximum_segments, + ) if data is None: return [] return [{"d": data, "fill": colour, "fill-rule": "evenodd"}] @@ -1384,13 +1674,17 @@ def generate_svg( masks: list[np.ndarray] | None = None, *, min_pixels: int = 32, - min_impact: float = 1e-5, + min_impact: float = 3e-6, max_layers: int = 512, segments: int = 16, + curvature_threshold: float | None = None, + maximum_segments: int = 2048, fill_holes: bool = True, hybrid_strokes: bool = True, ocr: bool = True, max_side: int | None = SAMVG_MAX_SIDE, + model: str = SAMVG_MODEL, + points_per_batch: int = SAMVG_POINTS_PER_BATCH, rasterize: Callable[[str, int, int], bytes] | None = None, ) -> str: """Generate SAMVG's traced, pre-optimisation SVG from a target image.""" @@ -1412,12 +1706,23 @@ def generate_svg( max_layers=max_layers, fill_holes=fill_holes, max_side=max_side, + model=model, + points_per_batch=points_per_batch, ) ) + # ``retrieve_layers`` has already done this for the normal SAM path. Do + # it here too for caller-supplied masks, which otherwise would export + # broad lower fills coloured by pixels that later paths hide. + if masks is not None: + layers = recolour_visible_layers(image, layers) paths = [] for layer in layers: for attributes in _layer_svg_attributes( - layer, segments, hybrid_strokes=hybrid_strokes + layer, + segments, + hybrid_strokes=hybrid_strokes, + curvature_threshold=curvature_threshold, + maximum_segments=maximum_segments, ): markup = " ".join(f'{key}="{value}"' for key, value in attributes.items()) paths.append(f"") @@ -1436,14 +1741,11 @@ def residual_prompt_points( target: Image.Image, rendered: Image.Image, *, - radius_fraction: float = 0.06, + radius_fraction: float = SAMVG_RESIDUAL_RADIUS_FRACTION, threshold: float = 0.784, - max_points: int = 16, + max_points: int | None = None, ) -> list[tuple[int, int]]: """Locate SAMVG's convolved, thresholded residual components.""" - import torch - import torch.nn.functional as functional - target_pixels = np.asarray(target.convert("RGB"), dtype=np.float32) / 255.0 rendered_pixels = np.asarray(rendered.convert("RGB"), dtype=np.float32) / 255.0 # SAMVG sums RGB-channel difference before applying its 0.784 threshold. @@ -1451,24 +1753,9 @@ def residual_prompt_points( difference = np.abs(target_pixels - rendered_pixels).sum(axis=2) height, width = difference.shape radius = max(2, round(min(height, width) * radius_fraction)) - yy, xx = np.ogrid[-radius : radius + 1, -radius : radius + 1] - kernel = (xx * xx + yy * yy <= radius * radius).astype(np.float32) - # Reflected padding preserves the prior symmetric-boundary definition; - # FFT convolution keeps the full-resolution recovery pass practical. - padded = np.pad(difference, radius, mode="symmetric") - smoothed = functional.conv2d( - torch.from_numpy(padded)[None, None], - torch.from_numpy((kernel / kernel.sum())[None, None]), - )[0, 0].numpy() - labels, count = _label(smoothed >= threshold) - points: list[tuple[float, int, int]] = [] - for index in range(1, count + 1): - ys, xs = np.nonzero(labels == index) - if len(xs): - points.append( - (float(smoothed[ys, xs].mean()), round(xs.mean()), round(ys.mean())) - ) - return [(x, y) for _score, x, y in sorted(points, reverse=True)[:max_points]] + return _circular_component_centres( + difference, radius, threshold=threshold, max_points=max_points + ) def _append_layers( @@ -1477,12 +1764,18 @@ def _append_layers( segments: int, *, hybrid_strokes: bool = True, + curvature_threshold: float | None = None, + maximum_segments: int = 2048, ) -> str: """Add newly prompted paths to an already optimised SVG.""" root = ET.fromstring(svg) for layer in layers: for attributes in _layer_svg_attributes( - layer, segments, hybrid_strokes=hybrid_strokes + layer, + segments, + hybrid_strokes=hybrid_strokes, + curvature_threshold=curvature_threshold, + maximum_segments=maximum_segments, ): ET.SubElement( root, @@ -1556,13 +1849,15 @@ def _accept_text_layers( def _accepted_fit( - svg: str, image: Image.Image, *, rasterize, steps: int + svg: str, image: Image.Image, *, rasterize, steps: int, learn_alpha: bool = False ) -> tuple[str, Image.Image]: """Keep a differentiable fit only when the actual SVG renderer improves.""" from vectrify.refine.paths import fit_filled_svg_bounded before = _render_svg(svg, image, rasterize) - fitted = fit_filled_svg_bounded(svg, image, rasterize=rasterize, steps=steps) + fitted = fit_filled_svg_bounded( + svg, image, rasterize=rasterize, steps=steps, learn_alpha=learn_alpha + ) after = _render_svg(fitted, image, rasterize) if _mse(image, after) <= _mse(image, before): return fitted, after @@ -1576,16 +1871,21 @@ def vectorize_svg( rasterize, steps: int = 500, min_pixels: int = 32, - min_impact: float = 1e-5, + min_impact: float = 3e-6, max_layers: int = 512, segments: int = 16, max_side: int | None = SAMVG_MAX_SIDE, + learn_alpha: bool = False, + curvature_threshold: float | None = None, + maximum_segments: int = 2048, ) -> str: """Run SAMVG's two 500-step optimise-and-recover phases. ``rasterize`` is the format backend's renderer, used solely to form the residual map after the first pass. The actual differentiable fit is the built-in filled-path optimiser so SAMVG has no external renderer dependency. + ``learn_alpha`` and ``curvature_threshold`` select the dissertation's + SAMVG+alpha and SAMVG+var representation variations, respectively. """ image = image.convert("RGB") runtime = _sam_runtime() @@ -1605,18 +1905,27 @@ def vectorize_svg( layers, segments, hybrid_strokes=False, + curvature_threshold=curvature_threshold, + maximum_segments=maximum_segments, ) first, first_render = _accepted_fit( - initial, image, rasterize=rasterize, steps=steps + initial, + image, + rasterize=rasterize, + steps=steps, + learn_alpha=learn_alpha, ) points = residual_prompt_points(image, first_render) - _canvas, coverage = _render_layers((image.height, image.width), layers) added = filter_by_impact( image, prompted_masks(image, points, max_side=max_side, _runtime=runtime), existing=layers, initial_canvas=np.asarray(first_render, dtype=np.uint8), - initial_coverage=coverage, + # Residual recovery scores against the fitted raster, not a blank + # segmentation canvas. Every pixel therefore has ordinary raster + # error; marking holes in the old masks uncovered would falsely + # reward any prompted mask placed there. + initial_coverage=np.ones((image.height, image.width), dtype=bool), min_pixels=min_pixels, min_impact=min_impact, max_layers=max_layers, @@ -1626,12 +1935,27 @@ def vectorize_svg( len(points), len(added), ) - return _accepted_fit( - _append_layers(first, added, segments, hybrid_strokes=False), + final, final_render = _accepted_fit( + _append_layers( + first, + added, + segments, + hybrid_strokes=False, + curvature_threshold=curvature_threshold, + maximum_segments=maximum_segments, + ), image, rasterize=rasterize, steps=steps, - )[0] + learn_alpha=learn_alpha, + ) + # A locally accepted second fit can still be worse than the first fit + # if its residual additions were harmful. The public two-phase result + # must never discard an already accepted Cairo-raster improvement. + if _mse(image, final_render) <= _mse(image, first_render): + return final + log.info("SAMVG residual phase rejected: it increased exported SVG MSE.") + return first finally: del runtime try: diff --git a/src/vectrify/vector/runner.py b/src/vectrify/vector/runner.py index a655975..9db5507 100644 --- a/src/vectrify/vector/runner.py +++ b/src/vectrify/vector/runner.py @@ -39,7 +39,12 @@ resize_long_side, ) from vectrify.llm.models import api_key_env -from vectrify.refine.samvg import generate_svg +from vectrify.refine.samvg import ( + SAMVG_MAX_SIDE, + SAMVG_MODEL, + SAMVG_POINTS_PER_BATCH, + generate_svg, +) from vectrify.score import ScorerType, choose_scorer from vectrify.score.base import DEFAULT_CONFIG from vectrify.score.compare import compare, prepare @@ -112,7 +117,17 @@ class VectorSearchConfig: vision_model: str = DEFAULT_VISION_MODEL auto_crop: bool = True segment_count: int = 8 - samvg_seed: bool = True + samvg_seed: bool = False + samvg_model: str = SAMVG_MODEL + samvg_max_side: int = SAMVG_MAX_SIDE + samvg_points_per_batch: int = SAMVG_POINTS_PER_BATCH + samvg_min_pixels: int = 32 + samvg_min_impact: float = 3e-6 + samvg_max_layers: int = 512 + samvg_segments: int = 16 + samvg_fill_holes: bool = True + samvg_hybrid_strokes: bool = True + samvg_ocr: bool = True dry_run: bool = False @@ -316,7 +331,17 @@ def run_vector_search( vision_model: str = DEFAULT_VISION_MODEL, # for the front evaluator auto_crop: bool = True, segment_count: int = 8, - samvg_seed: bool = True, + samvg_seed: bool = False, + samvg_model: str = SAMVG_MODEL, + samvg_max_side: int = SAMVG_MAX_SIDE, + samvg_points_per_batch: int = SAMVG_POINTS_PER_BATCH, + samvg_min_pixels: int = 32, + samvg_min_impact: float = 3e-6, + samvg_max_layers: int = 512, + samvg_segments: int = 16, + samvg_fill_holes: bool = True, + samvg_hybrid_strokes: bool = True, + samvg_ocr: bool = True, dry_run: bool = False, dry_run_parameters: Mapping[str, Any] | None = None, stats: "SearchStats | None" = None, @@ -351,6 +376,16 @@ def run_vector_search( auto_crop=auto_crop, segment_count=segment_count, samvg_seed=samvg_seed, + samvg_model=samvg_model, + samvg_max_side=samvg_max_side, + samvg_points_per_batch=samvg_points_per_batch, + samvg_min_pixels=samvg_min_pixels, + samvg_min_impact=samvg_min_impact, + samvg_max_layers=samvg_max_layers, + samvg_segments=samvg_segments, + samvg_fill_holes=samvg_fill_holes, + samvg_hybrid_strokes=samvg_hybrid_strokes, + samvg_ocr=samvg_ocr, dry_run=dry_run, ) resolution_llm = config.resolution_llm @@ -377,6 +412,16 @@ def run_vector_search( auto_crop = config.auto_crop segment_count = config.segment_count samvg_seed = config.samvg_seed + samvg_model = config.samvg_model + samvg_max_side = config.samvg_max_side + samvg_points_per_batch = config.samvg_points_per_batch + samvg_min_pixels = config.samvg_min_pixels + samvg_min_impact = config.samvg_min_impact + samvg_max_layers = config.samvg_max_layers + samvg_segments = config.samvg_segments + samvg_fill_holes = config.samvg_fill_holes + samvg_hybrid_strokes = config.samvg_hybrid_strokes + samvg_ocr = config.samvg_ocr dry_run = config.dry_run epoch_seeds = resolve_seeds(seeds) @@ -481,6 +526,16 @@ def run_vector_search( content = format_plugin.extract_from_llm( generate_svg( original_img, + min_pixels=samvg_min_pixels, + min_impact=samvg_min_impact, + max_layers=samvg_max_layers, + segments=samvg_segments, + fill_holes=samvg_fill_holes, + hybrid_strokes=samvg_hybrid_strokes, + ocr=samvg_ocr, + max_side=samvg_max_side, + model=samvg_model, + points_per_batch=samvg_points_per_batch, rasterize=lambda svg, width, height: format_plugin.rasterize( svg, out_w=width, out_h=height ), diff --git a/tests/refine/test_filled_paths.py b/tests/refine/test_filled_paths.py index 9428fcd..8fce507 100644 --- a/tests/refine/test_filled_paths.py +++ b/tests/refine/test_filled_paths.py @@ -16,6 +16,7 @@ _large_path_tile_candidates, _pad_fused_cubics, _tiled_large_path_coverage, + _torch_compile_enabled, _xing_loss, fit_filled_svg, fit_opaque_fills_locally, @@ -37,6 +38,11 @@ def _sixteen_cubic_circle(torch): )[None] +def test_torch_compile_can_be_explicitly_disabled(monkeypatch): + monkeypatch.setenv("TORCH_COMPILE_DISABLE", "1") + assert not _torch_compile_enabled() + + @pytest.mark.parametrize("samples", [8, 16, 32]) def test_native_winding_matches_torch_forward_and_gradient(samples, monkeypatch): """Release-wheel CUDA path agrees with the portable sampled renderer.""" @@ -261,6 +267,93 @@ def test_filled_path_fit_moves_fill_colour_toward_target(): assert _mse(fitted, target) < _mse(SVG, target) +def test_filled_path_fit_can_learn_fill_opacity(): + source = SVG.replace('#0000ff"', '#ff0000" fill-opacity="1"') + target = Image.new("RGB", (24, 24), "black") + target.paste("#400000", (4, 4, 16, 16)) + + fitted = fit_filled_svg( + source, + target, + steps=12, + point_learning_rate=0.0, + color_learning_rate=0.1, + learn_alpha=True, + ) + + root = ET.fromstring(fitted) + path = next(element for element in root.iter() if element.get("d")) + assert 0 < float(path.get("fill-opacity", "0")) < 1 + + +def test_sparse_fill_replay_matches_dense_replay_update(): + source = SVG.replace( + "", + '', + ) + target = Image.new("RGB", (24, 24), "black") + target.paste("red", (4, 4, 16, 16)) + + dense = fit_filled_svg( + source, + target, + steps=1, + point_learning_rate=0.0, + color_learning_rate=0.1, + ) + sparse = fit_filled_svg( + source, + target, + steps=1, + point_learning_rate=0.0, + color_learning_rate=0.1, + sparse_replay=True, + ) + + assert abs(_mse(dense, target) - _mse(sparse, target)) < 2 + + +def test_filled_path_fit_preserves_a_closed_contours_segment_count(): + target = Image.new("RGB", (24, 24), "black") + target.paste("red", (4, 4, 16, 16)) + + fitted = fit_filled_svg( + SVG, + target, + steps=1, + point_learning_rate=0.5, + color_learning_rate=0.0, + ) + + root = ET.fromstring(fitted) + path = next(element for element in root.iter() if element.get("d")) + assert [len(contour) for contour in parse_filled_cubics(path.get("d", ""))] == [4] + + +def test_filled_path_fit_preserves_each_compound_contours_closure(): + svg = ( + '' + '' + "" + ) + + fitted = fit_filled_svg( + svg, + Image.new("RGB", (24, 24), "black"), + steps=1, + point_learning_rate=0.5, + color_learning_rate=0.0, + ) + + root = ET.fromstring(fitted) + path = next(element for element in root.iter() if element.get("d")) + assert [len(contour) for contour in parse_filled_cubics(path.get("d", ""))] == [ + 4, + 4, + ] + + def test_local_fill_fit_changes_only_one_bounded_group(): paths = [] for index in range(20): @@ -631,6 +724,19 @@ def unexpected_sampled_fallback(*_args, **_kwargs): assert "path" in fitted +def test_filled_fit_initialises_a_variable_length_contour(): + commands = " ".join("C 0 0 1 0 2 0" for _ in range(17)) + svg = ( + '' + f'' + "" + ) + + fitted = fit_filled_svg(svg, Image.new("RGB", (8, 8), "white"), steps=0) + + assert "path" in fitted + + def test_bounded_compositing_gradient_matches_monolithic_render(): """The memory-bounded fit pass must retain the full painter's-order MSE gradient.""" torch = pytest.importorskip("torch") diff --git a/tests/refine/test_samvg.py b/tests/refine/test_samvg.py index fea618b..459ffb9 100644 --- a/tests/refine/test_samvg.py +++ b/tests/refine/test_samvg.py @@ -104,9 +104,10 @@ def generate(self, **kwargs): def test_accepted_fit_uses_bounded_fill_coordinate_descent(monkeypatch): seen = {} - def bounded(svg, image, *, rasterize, steps): + def bounded(svg, image, *, rasterize, steps, learn_alpha): seen["image"] = image.size seen["steps"] = steps + seen["learn_alpha"] = learn_alpha assert rasterize is not None return svg @@ -121,7 +122,7 @@ def bounded(svg, image, *, rasterize, steps): assert fitted == svg assert rendered.size == image.size - assert seen == {"image": (16, 16), "steps": 7} + assert seen == {"image": (16, 16), "steps": 7, "learn_alpha": False} def test_vectorize_svg_runs_a_second_residual_recovery_phase(monkeypatch): @@ -133,7 +134,7 @@ def test_vectorize_svg_runs_a_second_residual_recovery_phase(monkeypatch): initial_layer = MaskLayer(base, (10, 20, 30), 1.0) added_layer = MaskLayer(added, (40, 50, 60), 1.0) calls = [] - monkeypatch.setattr(samvg, "_sam_runtime", lambda: object()) + monkeypatch.setattr(samvg, "_sam_runtime", lambda **_kwargs: object()) monkeypatch.setattr( samvg, "retrieve_layers", lambda *_args, **_kwargs: [initial_layer] ) @@ -145,7 +146,7 @@ def test_vectorize_svg_runs_a_second_residual_recovery_phase(monkeypatch): lambda _image, _masks, **kwargs: [*kwargs["existing"], added_layer], ) - def accepted(svg, _image, *, rasterize, steps): + def accepted(svg, _image, *, rasterize, steps, learn_alpha): assert rasterize is not None calls.append( ( @@ -154,20 +155,56 @@ def accepted(svg, _image, *, rasterize, steps): for element in ET.fromstring(svg).iter() ), steps, + learn_alpha, ) ) return svg, image monkeypatch.setattr(samvg, "_accepted_fit", accepted) - result = samvg.vectorize_svg(image, rasterize=SvgPlugin().rasterize, steps=3) + result = samvg.vectorize_svg( + image, rasterize=SvgPlugin().rasterize, steps=3, learn_alpha=True + ) - assert calls == [(1, 3), (2, 3)] + assert calls == [(1, 3, True), (2, 3, True)] root = ET.fromstring(result) paths = list(root.findall("{http://www.w3.org/2000/svg}path")) assert len(paths) == 2 assert all(path.get("stroke") is None for path in paths) +def test_vectorize_svg_rejects_a_residual_phase_that_regresses_first_fit(monkeypatch): + image = Image.new("RGB", (16, 16), "white") + base = np.zeros((16, 16), dtype=bool) + base[2:10, 2:10] = True + added = np.zeros((16, 16), dtype=bool) + added[10:14, 10:14] = True + initial_layer = MaskLayer(base, (10, 20, 30), 1.0) + added_layer = MaskLayer(added, (40, 50, 60), 1.0) + monkeypatch.setattr(samvg, "_sam_runtime", lambda **_kwargs: object()) + monkeypatch.setattr( + samvg, "retrieve_layers", lambda *_args, **_kwargs: [initial_layer] + ) + monkeypatch.setattr(samvg, "residual_prompt_points", lambda *_args: [(12, 12)]) + monkeypatch.setattr(samvg, "prompted_masks", lambda *_args, **_kwargs: [added]) + monkeypatch.setattr( + samvg, + "filter_by_impact", + lambda _image, _masks, **kwargs: [*kwargs["existing"], added_layer], + ) + renders = [image, Image.new("RGB", image.size, "black")] + monkeypatch.setattr( + samvg, + "_accepted_fit", + lambda svg, _image, **_kwargs: (svg, renders.pop(0)), + ) + + result = samvg.vectorize_svg(image, rasterize=SvgPlugin().rasterize, steps=3) + + root = ET.fromstring(result) + path_count = sum(element.tag.endswith("path") for element in root.iter()) + assert path_count == 1 + + def test_generate_svg_writes_detected_words_as_editable_text(monkeypatch): monkeypatch.setattr(samvg, "retrieve_layers", lambda *_args, **_kwargs: []) monkeypatch.setattr( @@ -265,6 +302,30 @@ def __call__(self, source, **kwargs): assert all(mask.shape == (8, 12) for mask in masks) +def test_automatic_mask_finalization_suppresses_across_crop_sources(): + import torch + + masks = torch.tensor( + [ + [[True, True], [True, True]], + [[False, False], [False, True]], + ] + ) + # These candidates represent the same image-space crop box. The second + # candidate has a different raster but a higher SAM IoU, so AMG's one + # image-global NMS must retain it instead of allowing each source to keep + # its own duplicate. + retained = samvg._finalize_automatic_masks( + masks, + torch.tensor([0.8, 0.9]), + torch.tensor([[0.0, 0.0, 2.0, 2.0], [0.0, 0.0, 2.0, 2.0]]), + ) + + assert len(retained) == 1 + assert retained[0][1, 1] + assert not retained[0][0, 0] + + def test_retrieve_layers_reuses_one_runtime_for_automatic_and_coverage_prompts( monkeypatch, ): @@ -307,6 +368,40 @@ def test_filter_by_impact_keeps_useful_nested_masks_in_layer_order(): assert all(layer.impact > 0 for layer in layers) +def test_filter_by_impact_scores_a_disconnected_mask_before_emitting_components(): + pixels = np.zeros((12, 12, 3), dtype=np.uint8) + pixels[2:5, 2:5] = (220, 20, 20) + pixels[7:10, 7:10] = (20, 20, 220) + image = Image.fromarray(pixels) + mask = np.zeros((12, 12), dtype=bool) + mask[2:5, 2:5] = True + mask[7:10, 7:10] = True + + layers = filter_by_impact(image, [mask], min_pixels=1, min_impact=0) + + assert [int(layer.mask.sum()) for layer in layers] == [9, 9] + assert {layer.colour for layer in layers} == {(120, 20, 120)} + assert layers[0].impact == layers[1].impact + + +def test_filter_by_impact_residual_canvas_does_not_charge_covered_pixels_as_blank(): + image = Image.new("RGB", (8, 8), (128, 128, 128)) + mask = np.zeros((8, 8), dtype=bool) + mask[2:6, 2:6] = True + fitted = np.full((8, 8, 3), 128, dtype=np.uint8) + + layers = filter_by_impact( + image, + [mask], + initial_canvas=fitted, + initial_coverage=np.ones((8, 8), dtype=bool), + min_pixels=1, + min_impact=1e-6, + ) + + assert layers == [] + + def test_incremental_impact_scoring_matches_full_canvas_recomputation(): pixels = np.full((16, 16, 3), 255, dtype=np.uint8) pixels[2:12, 2:12] = (180, 60, 30) @@ -397,6 +492,44 @@ def test_components_fill_only_tiny_enclosed_holes(): assert components[0].all() +def test_bounded_component_hole_checks_match_full_canvas_semantics(): + def full_canvas(mask: np.ndarray, min_pixels: int) -> list[np.ndarray]: + result = [] + for runs in samvg._run_components(mask): + if sum(end - start for _y, start, end in runs) < min_pixels: + continue + component = np.zeros(mask.shape, dtype=bool) + for y, start, end in runs: + component[y, start:end] = True + for hole in samvg._run_components(~component): + area = sum(end - start for _y, start, end in hole) + touches_border = any( + y in {0, mask.shape[0] - 1} or start == 0 or end == mask.shape[1] + for y, start, end in hole + ) + if area <= min_pixels and not touches_border: + for y, start, end in hole: + component[y, start:end] = True + result.append(component) + return result + + mask = np.zeros((20, 24), dtype=bool) + mask[1:12, 1:12] = True + mask[4:6, 4:6] = False + mask[7:10, 7:10] = False + mask[3:5, 18:20] = True + mask[14:17, 2:5] = True + + bounded = _components(mask, min_pixels=4) + original = full_canvas(mask, min_pixels=4) + + assert len(bounded) == len(original) + assert all( + np.array_equal(left, right) + for left, right in zip(bounded, original, strict=True) + ) + + def test_internal_morphology_matches_scipy_default_connectivity(): mask = np.array( [[False, True, True], [True, True, True], [True, True, True]], dtype=bool @@ -432,6 +565,30 @@ def test_mask_path_keeps_a_hole_as_a_second_even_odd_subpath(): assert path.count(" Z") == 2 +def test_mask_path_supports_the_variable_segment_tracing_variation(): + mask = np.zeros((48, 48), dtype=bool) + mask[8:40, 8:40] = True + mask[16:32, 16:32] = False + + path = mask_path(mask, curvature_threshold=0.8, maximum_segments=6) + + assert path is not None + assert path.count("M ") == 2 + assert 6 <= path.count("C ") <= 12 + + +def test_variable_corners_retains_nearby_local_extrema(monkeypatch): + monkeypatch.setattr( + samvg, + "_curvature_scores", + lambda _loop: np.array((1.0, -0.9, 1.0, -0.8, 1.0, -0.7, 1.0, -0.6)), + ) + + corners = samvg._variable_corners([(0.0, 0.0)] * 8, threshold=0, maximum=16) + + assert corners == [1, 3, 5, 7] + + def test_generate_svg_creates_editable_layered_paths_from_supplied_masks(): image = Image.new("RGB", (10, 8), "white") pixels = np.asarray(image).copy() @@ -449,6 +606,28 @@ def test_generate_svg_creates_editable_layered_paths_from_supplied_masks(): assert paths[0].get("fill") == "#1482dc" +def test_generate_svg_refits_each_visible_fill_colour_after_mask_selection(): + pixels = np.full((12, 12, 3), (220, 30, 30), dtype=np.uint8) + pixels[4:8, 4:8] = (20, 40, 230) + image = Image.fromarray(pixels) + outer = np.ones((12, 12), dtype=bool) + inner = np.zeros((12, 12), dtype=bool) + inner[4:8, 4:8] = True + + root = ET.fromstring( + generate_svg( + image, + [outer, inner], + min_pixels=1, + min_impact=0, + ocr=False, + ) + ) + paths = list(root.findall("{http://www.w3.org/2000/svg}path")) + + assert [path.get("fill") for path in paths] == ["#dc1e1e", "#1428e6"] + + def test_thin_single_contour_mask_is_emitted_as_a_round_stroke(): image = Image.new("RGB", (12, 32), "white") pixels = np.asarray(image).copy() @@ -519,6 +698,7 @@ def test_coverage_prompt_points_selects_the_centre_of_a_large_empty_region(): [MaskLayer(occupied, (10, 20, 30), 1.0)], (32, 32), radius_fraction=0.15, + max_points=10, ) assert points @@ -569,6 +749,32 @@ def test_cubic_fit_reparameterises_nonuniform_curve_samples(): ) +def test_fixed_cubic_tracing_does_not_duplicate_each_curve_endpoint(monkeypatch): + loop = [ + (0.0, 0.0), + (1.0, 0.0), + (2.0, 0.0), + (2.0, 1.0), + (2.0, 2.0), + (1.0, 2.0), + (0.0, 2.0), + (0.0, 1.0), + ] + samples = [] + monkeypatch.setattr(samvg, "_corners", lambda *_args: [0, 2, 4, 6]) + + def fit(points): + samples.append(points.copy()) + return points[0], points[-1] + + monkeypatch.setattr(samvg, "_fit_cubic", fit) + + samvg._cubic_loop(loop, segments=4) + + assert [len(points) for points in samples] == [3, 3, 3, 3] + assert all(not np.array_equal(points[-1], points[-2]) for points in samples) + + def test_sam_input_cap_restores_binary_masks_to_the_original_canvas(): image = Image.new("RGB", (100, 50), "white") diff --git a/tests/test_cli.py b/tests/test_cli.py index 82f30e0..6713f87 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -61,10 +61,9 @@ def test_negative_seeds_raises(): parse_args(["img.png", "--seeds", "-1"]) -def test_seeds_zero_uses_the_default_samvg_seed_or_requires_resume_when_disabled(): - assert parse_args(["img.png", "--seeds", "0"]).seeds == 0 +def test_seeds_zero_requires_resume_or_an_explicit_samvg_seed(): with pytest.raises(SystemExit): - parse_args(["img.png", "--seeds", "0", "--no-samvg-seed"]) + parse_args(["img.png", "--seeds", "0"]) assert parse_args(["img.png", "--seeds", "0", "--resume"]).seeds == 0 @@ -82,9 +81,46 @@ def test_samvg_seed_allows_a_local_only_run(): assert args.seeds == 0 -def test_samvg_seed_is_enabled_by_default_and_can_be_disabled(): - assert parse_args(["img.png"]).samvg_seed is True - assert parse_args(["img.png", "--no-samvg-seed"]).samvg_seed is False +def test_samvg_seed_is_disabled_by_default_and_can_be_enabled(): + assert parse_args(["img.png"]).samvg_seed is False + assert parse_args(["img.png", "--samvg-seed"]).samvg_seed is True + + +def test_samvg_seed_knobs_are_parsed_independently_of_the_opt_in_flag(): + args = parse_args( + [ + "img.png", + "--samvg-model", + "facebook/sam-vit-base", + "--samvg-max-side", + "768", + "--samvg-points-per-batch", + "96", + "--samvg-min-pixels", + "48", + "--samvg-min-impact", + "0.00002", + "--samvg-max-layers", + "128", + "--samvg-segments", + "12", + "--no-samvg-fill-holes", + "--no-samvg-hybrid-strokes", + "--no-samvg-ocr", + ] + ) + + assert args.samvg_seed is False + assert args.samvg_model == "facebook/sam-vit-base" + assert args.samvg_max_side == 768 + assert args.samvg_points_per_batch == 96 + assert args.samvg_min_pixels == 48 + assert args.samvg_min_impact == 0.00002 + assert args.samvg_max_layers == 128 + assert args.samvg_segments == 12 + assert args.samvg_fill_holes is False + assert args.samvg_hybrid_strokes is False + assert args.samvg_ocr is False # Defaults are pinned as literals so a change to any default is a visible,