diff --git a/README.md b/README.md
index 5a5bbeeb..93c2a0c8 100644
--- a/README.md
+++ b/README.md
@@ -118,6 +118,29 @@ 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.
+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.
+Set `VECTRIFY_SAMVG_MAX_SIDE` to choose another cap, or pass `max_side=None` to
+the Python API to opt out explicitly.
+
+Automatic SAM masks retain the dissertation's 32×32 prompt grid but decode 64
+prompts per CUDA batch in FP16 by default. Set `VECTRIFY_SAMVG_POINTS_PER_BATCH`
+for a larger-memory GPU; full-resolution mask filtering remains on CPU so the
+batch does not consume the renderer's CUDA memory.
+
The default candidate is deliberately the segmentation-and-tracing seed only;
it does not run the path optimiser. This keeps seed quality measurable without
mixing in local refinement.
+
+For the dissertation-style two-phase measurement (initial fit, residual prompts,
+and recovery fit), build with the optional native CUDA renderer and run:
+
+```sh
+VECTRIFY_BUILD_SAMVG_CUDA=1 uv build --wheel --no-build-isolation
+uv pip install --force-reinstall --no-deps dist/vectrify-*.whl
+.venv/bin/python scripts/bench_samvg_two_phase.py --cat
+```
+
+`--all` also evaluates every benchmark target and the connect-the-dots duck.
+Each target directory contains the five stage rasters, SVGs, a gallery,
+pixel-error table, and per-bounded-group CUDA memory/timing data.
diff --git a/pyproject.toml b/pyproject.toml
index d0e9d192..2eec9b63 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -55,17 +55,15 @@ vision = [
"scikit-learn>=1.3.0",
"torch>=2.0.0",
"torchvision>=0.28.0",
- "transformers>=4.40.0",
+ "transformers>=4.49.0",
]
# The CUDA extension is shipped in platform-specific wheels. Installing this
# extra is deliberately sufficient for the SAMVG seed even on machines that
# use the portable Torch renderer fallback.
samvg = [
- "scipy>=1.11.0",
- "scikit-learn>=1.3.0",
"torch>=2.0.0",
"torchvision>=0.28.0",
- "transformers>=4.40.0",
+ "transformers>=4.49.0",
]
graphviz = [
"graphviz>=0.21",
diff --git a/scripts/bench_samvg_two_phase.py b/scripts/bench_samvg_two_phase.py
new file mode 100644
index 00000000..69bb8c48
--- /dev/null
+++ b/scripts/bench_samvg_two_phase.py
@@ -0,0 +1,221 @@
+"""Run and record SAMVG's two-phase segmentation, fit, and recovery process.
+
+Examples:
+ uv run python scripts/bench_samvg_two_phase.py --cat
+ uv run python scripts/bench_samvg_two_phase.py --all
+
+The native CUDA extension must be available for the 1024px cat workload. The
+script deliberately uses the regular SAMVG masks and fixed-16-segment tracer;
+it only bounds the differentiable fit to one spatial fill group at a time.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from time import perf_counter
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from vectrify.formats.svg.plugin import SvgPlugin
+from vectrify.refine.paths import fit_filled_svg_bounded
+from vectrify.refine.samvg import (
+ _append_layers,
+ _mse,
+ _render_layers,
+ _render_svg,
+ filter_by_impact,
+ prompted_masks,
+ residual_prompt_points,
+ retrieve_layers,
+)
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _path_count(svg: str) -> int:
+ return sum(
+ element.tag.split("}")[-1] == "path" for element in ET.fromstring(svg).iter()
+ )
+
+
+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 _write_gallery(images: list[tuple[str, Image.Image]], destination: Path) -> None:
+ width = max(image.width for _name, image in images)
+ height = max(image.height for _name, image in images)
+ gallery = Image.new("RGB", (width * len(images), height + 28), "white")
+ labels = ImageDraw.Draw(gallery)
+ for index, (name, image) in enumerate(images):
+ gallery.paste(image.convert("RGB"), (index * width, 28))
+ labels.text((index * width + 4, 6), name, fill="black")
+ gallery.save(destination)
+
+
+def _fit_if_improved(
+ svg: str,
+ target: Image.Image,
+ plugin: SvgPlugin,
+ steps: int,
+) -> tuple[str, Image.Image, list[dict[str, int | float]], bool]:
+ before = _render_svg(svg, target, plugin.rasterize)
+ measurements: list[dict[str, int | float]] = []
+ candidate = fit_filled_svg_bounded(
+ svg,
+ target,
+ rasterize=plugin.rasterize,
+ steps=steps,
+ measurements=measurements,
+ )
+ after = _render_svg(candidate, target, plugin.rasterize)
+ if _mse(target, after) <= _mse(target, before):
+ return candidate, after, measurements, True
+ return svg, before, measurements, False
+
+
+def run_target(
+ target_path: Path,
+ output: Path,
+ *,
+ steps: int,
+ reference_svg: Path | None = None,
+) -> None:
+ target = Image.open(target_path).convert("RGB")
+ plugin = SvgPlugin()
+ destination = output / target_path.stem
+ destination.mkdir(parents=True, exist_ok=True)
+ started = perf_counter()
+ layers = retrieve_layers(target)
+ initial = _append_layers(
+ f'',
+ layers,
+ 16,
+ hybrid_strokes=False,
+ )
+ first, first_render, first_measurements, first_accepted = _fit_if_improved(
+ initial, target, plugin, steps
+ )
+ 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),
+ existing=layers,
+ initial_canvas=np.asarray(first_render, dtype=np.uint8),
+ initial_coverage=coverage,
+ )[len(layers) :]
+ recovery = _append_layers(first, added, 16, hybrid_strokes=False)
+ final, final_render, final_measurements, final_accepted = _fit_if_improved(
+ recovery, target, plugin, steps
+ )
+ stages = [
+ ("target", target, None),
+ ("first-seed", _render_svg(initial, target, plugin.rasterize), initial),
+ ("first-fit", first_render, first),
+ (
+ "residual-recovery",
+ _render_svg(recovery, target, plugin.rasterize),
+ recovery,
+ ),
+ ("final-fit", final_render, final),
+ ]
+ 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 = []
+ for name, rendered, svg in stages:
+ rendered.save(destination / f"{name}.png")
+ if svg is not None:
+ (destination / f"{name}.svg").write_text(svg)
+ rows.append(
+ {
+ "stage": name,
+ "l1": _l1(target, rendered),
+ "mse": _mse(target, rendered),
+ "paths": _path_count(svg) if svg is not None else 0,
+ }
+ )
+ _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", "l1", "mse", "paths"])
+ writer.writeheader()
+ writer.writerows(rows)
+ measurements = [
+ {"phase": "first-fit", **measurement} for measurement in first_measurements
+ ] + [{"phase": "final-fit", **measurement} for measurement in final_measurements]
+ (destination / "fit-groups.json").write_text(json.dumps(measurements, indent=2))
+ (destination / "summary.json").write_text(
+ json.dumps(
+ {
+ "target": str(target_path),
+ "first_fit_accepted": first_accepted,
+ "final_fit_accepted": final_accepted,
+ "initial_layers": len(layers),
+ "residual_layers": len(added),
+ "wall_seconds": perf_counter() - started,
+ "stages": rows,
+ },
+ indent=2,
+ )
+ )
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--target", type=Path, action="append", default=[])
+ parser.add_argument(
+ "--reference-svg",
+ type=Path,
+ help="Reference SVG to Cairo-rasterize alongside a single target.",
+ )
+ parser.add_argument("--cat", action="store_true")
+ parser.add_argument(
+ "--all", action="store_true", help="Run cat, duck, and all bench targets."
+ )
+ parser.add_argument(
+ "--output", type=Path, default=ROOT / "bench/results/samvg-two-phase"
+ )
+ parser.add_argument("--steps", type=int, default=500)
+ args = parser.parse_args()
+ targets = list(args.target)
+ if args.cat or args.all:
+ targets.append(Path("/tmp/SAMVG_thesis/cat1024.jpg"))
+ if args.all:
+ targets.extend(sorted((ROOT / "bench/cases").glob("*/target.png")))
+ targets.append(ROOT / "connect-the-dots-little-duck.png")
+ if not targets:
+ parser.error("give --target, --cat, or --all")
+ if args.steps < 1:
+ parser.error("--steps must be positive")
+ if args.reference_svg is not None and len(targets) != 1:
+ parser.error("--reference-svg requires exactly one target")
+ for target in targets:
+ reference_svg = args.reference_svg
+ if target == Path("/tmp/SAMVG_thesis/cat1024.jpg"):
+ candidate = target.with_suffix(".svg")
+ if candidate.exists():
+ reference_svg = candidate
+ run_target(
+ target,
+ args.output,
+ steps=args.steps,
+ reference_svg=reference_svg,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/vectrify/cli.py b/src/vectrify/cli.py
index 4ece47d4..b61c0fd6 100644
--- a/src/vectrify/cli.py
+++ b/src/vectrify/cli.py
@@ -201,7 +201,8 @@ def parse_args(args: list[str] | None = None) -> argparse.Namespace:
action=argparse.BooleanOptionalAction,
default=True,
help="Add one SAMVG-inspired SVG seed made from automatic SAM masks, "
- "impact filtering, and contour tracing. Requires vectrify[vision] and "
+ "impact filtering, contour tracing, and Torch OCR. Requires "
+ "vectrify[samvg] and "
"is available for SVG output only. Default: on",
)
diff --git a/src/vectrify/formats/svg/plugin.py b/src/vectrify/formats/svg/plugin.py
index 7b0cd66e..aab748ab 100644
--- a/src/vectrify/formats/svg/plugin.py
+++ b/src/vectrify/formats/svg/plugin.py
@@ -29,7 +29,9 @@
PATH_FIT,
UnsupportedPathError,
fit_available,
- fit_random_group,
+ fit_svg_primitives_locally,
+ fittable_opaque_fills,
+ fittable_strokes,
)
log = logging.getLogger(__name__)
@@ -157,16 +159,18 @@ def mutate(
if reference_png is None or not fit_available():
return content, PATH_FIT
try:
- return (
- fit_random_group(
- content,
- reference_png,
- rasterize=lambda svg, w, h: self.rasterize(svg, w, h),
- weights=targets,
- gpu_gate=self.gpu_gate,
- ),
- PATH_FIT,
- )
+ if fittable_opaque_fills(content) or fittable_strokes(content):
+ return (
+ fit_svg_primitives_locally(
+ content,
+ reference_png,
+ rasterize=lambda svg, w, h: self.rasterize(svg, w, h),
+ weights=targets,
+ gpu_gate=self.gpu_gate,
+ ),
+ PATH_FIT,
+ )
+ raise UnsupportedPathError("no supported SVG primitive to fit")
except UnsupportedPathError as exc:
log.debug(f"Nothing to fit: {exc}")
return content, PATH_FIT
diff --git a/src/vectrify/refine/_samvg_cuda.cu b/src/vectrify/refine/_samvg_cuda.cu
index f0739b3a..847d77ed 100644
--- a/src/vectrify/refine/_samvg_cuda.cu
+++ b/src/vectrify/refine/_samvg_cuda.cu
@@ -232,6 +232,128 @@ __global__ void coverage_backward_kernel(const float* controls, const float* ups
}
}
+// The stroke is the union of discs centred along its cubics. Clamping the
+// closest-point parameter to [0, 1] makes the end discs round caps, while the
+// minimum over adjacent cubics gives round joins without separate join code.
+__device__ inline void closest_stroke_point(const float* path, float px, float py,
+ int& best_cubic, float& best_t,
+ float& best_x, float& best_y,
+ float& best_distance_sq) {
+ best_distance_sq = 1e30f;
+ best_cubic = 0; best_t = 0.f; best_x = path[0]; best_y = path[1];
+ for (int cubic = 0; cubic < kCubics; ++cubic) {
+ for (int seed = 0; seed < 5; ++seed) {
+ float t = .25f * seed;
+ for (int iteration = 0; iteration < 3; ++iteration) {
+ const float qx = cubic_component(path, cubic, t, 0) - px;
+ const float qy = cubic_component(path, cubic, t, 1) - py;
+ const float dx = cubic_derivative(path, cubic, t, 0);
+ const float dy = cubic_derivative(path, cubic, t, 1);
+ t = fminf(1.f, fmaxf(0.f, t - (qx*dx + qy*dy) / (dx*dx + dy*dy + 1e-6f)));
+ }
+ const float qx = cubic_component(path, cubic, t, 0);
+ const float qy = cubic_component(path, cubic, t, 1);
+ const float dx = qx - px, dy = qy - py;
+ const float distance_sq = dx*dx + dy*dy;
+ if (distance_sq < best_distance_sq) {
+ best_distance_sq = distance_sq; best_cubic = cubic; best_t = t;
+ best_x = qx; best_y = qy;
+ }
+ }
+ }
+}
+
+__global__ void stroke_forward_kernel(const float* controls, const float* widths,
+ float* output, int batches, int height, int width,
+ int subpixels, float x_base, float y_base) {
+ const int pixels = height * width, batch = blockIdx.x;
+ if (batch >= batches) return;
+ const float* path = controls + batch * kCubics * 8;
+ const float radius = fmaxf(0.f, widths[batch]) * .5f;
+ __shared__ float bounds[4];
+ path_bounds(path, bounds);
+ for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) {
+ const float base_x = x_base + float(pixel % width);
+ const float base_y = y_base + float(pixel / width);
+ float covered = 0.f;
+ if (base_x >= bounds[0] - radius - 2.f && base_x <= bounds[2] + radius + 2.f &&
+ base_y >= bounds[1] - radius - 2.f && base_y <= bounds[3] + radius + 2.f) {
+ for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) {
+ const float px = base_x + (float(subpixel % subpixels) + .5f) / subpixels;
+ const float py = base_y + (float(subpixel / subpixels) + .5f) / subpixels;
+ int cubic; float t, qx, qy, distance_sq;
+ closest_stroke_point(path, px, py, cubic, t, qx, qy, distance_sq);
+ const float distance = sqrtf(distance_sq + 1e-12f);
+ covered += 1.f / (1.f + expf((distance - radius) / .25f));
+ }
+ }
+ output[batch * pixels + pixel] = covered / float(subpixels * subpixels);
+ }
+}
+
+__global__ void stroke_backward_kernel(const float* controls, const float* widths,
+ const float* upstream, float* gradients,
+ float* width_gradients, int batches, int height,
+ int width, int subpixels, float x_base, float y_base) {
+ const int pixels = height * width, batch = blockIdx.x;
+ if (batch >= batches) return;
+ const float* path = controls + batch * kCubics * 8;
+ float* gradient = gradients + batch * kCubics * 8;
+ const float radius = fmaxf(0.f, widths[batch]) * .5f;
+ __shared__ float bounds[4];
+ __shared__ float reduction[8][256];
+ __shared__ float width_reduction[256];
+ float accumulated[kCubics * 8] = {0.f};
+ float accumulated_width = 0.f;
+ path_bounds(path, bounds);
+ for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) {
+ const float base_x = x_base + float(pixel % width);
+ const float base_y = y_base + float(pixel / width);
+ if (base_x < bounds[0] - radius - 2.f || base_x > bounds[2] + radius + 2.f ||
+ base_y < bounds[1] - radius - 2.f || base_y > bounds[3] + radius + 2.f) continue;
+ const float d_output = upstream[batch * pixels + pixel] / float(subpixels * subpixels);
+ for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) {
+ const float px = base_x + (float(subpixel % subpixels) + .5f) / subpixels;
+ const float py = base_y + (float(subpixel / subpixels) + .5f) / subpixels;
+ int cubic; float t, qx, qy, distance_sq;
+ closest_stroke_point(path, px, py, cubic, t, qx, qy, distance_sq);
+ const float distance = sqrtf(distance_sq + 1e-12f);
+ const float alpha = 1.f / (1.f + expf((distance - radius) / .25f));
+ const float edge = d_output * alpha * (1.f - alpha) / .25f;
+ const float u = 1.f - t;
+ const float basis[4] = {u*u*u, 3.f*u*u*t, 3.f*u*t*t, t*t*t};
+ for (int control = 0; control < 4; ++control) {
+ const int offset = cubic * 8 + control * 2;
+ accumulated[offset] -= edge * (qx - px) * basis[control] / distance;
+ accumulated[offset + 1] -= edge * (qy - py) * basis[control] / distance;
+ }
+ accumulated_width += edge * .5f;
+ }
+ }
+ for (int cubic = 0; cubic < kCubics; ++cubic) {
+ for (int component = 0; component < 8; ++component)
+ reduction[component][threadIdx.x] = accumulated[cubic * 8 + component];
+ __syncthreads();
+ for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+ if (threadIdx.x < stride)
+ for (int component = 0; component < 8; ++component)
+ reduction[component][threadIdx.x] += reduction[component][threadIdx.x + stride];
+ __syncthreads();
+ }
+ if (threadIdx.x == 0)
+ for (int component = 0; component < 8; ++component)
+ gradient[cubic * 8 + component] = reduction[component][0];
+ __syncthreads();
+ }
+ width_reduction[threadIdx.x] = accumulated_width;
+ __syncthreads();
+ for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
+ if (threadIdx.x < stride) width_reduction[threadIdx.x] += width_reduction[threadIdx.x + stride];
+ __syncthreads();
+ }
+ if (threadIdx.x == 0) width_gradients[batch] = width_reduction[0];
+}
+
__device__ inline void contours_bounds(const float* controls, int first, int last, float* bounds) {
if (threadIdx.x == 0) {
const float* first_path = controls + first * kCubics * 8;
@@ -617,6 +739,37 @@ torch::Tensor coverage_backward(torch::Tensor controls, torch::Tensor upstream,
return gradients;
}
+torch::Tensor stroke_forward(torch::Tensor controls, torch::Tensor widths, int64_t height,
+ int64_t width, int64_t subpixels, double x_origin,
+ double y_origin) {
+ TORCH_CHECK(controls.is_cuda() && controls.scalar_type() == torch::kFloat32);
+ TORCH_CHECK(widths.is_cuda() && widths.scalar_type() == torch::kFloat32);
+ TORCH_CHECK(widths.dim() == 1 && widths.size(0) == controls.size(0));
+ TORCH_CHECK(subpixels >= 1 && subpixels <= 4);
+ at::cuda::CUDAGuard guard(controls.device());
+ auto output = torch::zeros({controls.size(0), height, width}, controls.options());
+ stroke_forward_kernel<<>>(
+ controls.data_ptr(), widths.data_ptr(), output.data_ptr(),
+ controls.size(0), height, width, subpixels, float(x_origin), float(y_origin));
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+ return output;
+}
+
+std::vector stroke_backward(torch::Tensor controls, torch::Tensor widths,
+ torch::Tensor upstream, int64_t height,
+ int64_t width, int64_t subpixels,
+ double x_origin, double y_origin) {
+ at::cuda::CUDAGuard guard(controls.device());
+ auto gradients = torch::zeros_like(controls);
+ auto width_gradients = torch::zeros_like(widths);
+ stroke_backward_kernel<<>>(
+ controls.data_ptr(), widths.data_ptr(), upstream.data_ptr(),
+ gradients.data_ptr(), width_gradients.data_ptr(), controls.size(0),
+ height, width, subpixels, float(x_origin), float(y_origin));
+ C10_CUDA_KERNEL_LAUNCH_CHECK();
+ return {gradients, width_gradients};
+}
+
torch::Tensor multi_coverage_forward(torch::Tensor controls, torch::Tensor offsets,
int64_t height, int64_t width, int64_t subpixels,
double x_origin, double y_origin, bool evenodd) {
@@ -693,6 +846,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("backwards", &backwards);
m.def("coverage_forward", &coverage_forward);
m.def("coverage_backward", &coverage_backward);
+ m.def("stroke_forward", &stroke_forward);
+ m.def("stroke_backward", &stroke_backward);
m.def("multi_coverage_forward", &multi_coverage_forward);
m.def("multi_coverage_forward_topology", &multi_coverage_forward_topology);
m.def("multi_coverage_backward", &multi_coverage_backward);
diff --git a/src/vectrify/refine/cuda_renderer.py b/src/vectrify/refine/cuda_renderer.py
index 5cb25893..3dce9031 100644
--- a/src/vectrify/refine/cuda_renderer.py
+++ b/src/vectrify/refine/cuda_renderer.py
@@ -184,6 +184,62 @@ def backward(ctx: Any, *upstreams: Any) -> Any:
return Coverage.apply(controls)
+def stroke_coverage(
+ controls: Any,
+ widths: Any,
+ box: tuple[int, int, int, int],
+ *,
+ subpixels: int = 2,
+) -> Any | None:
+ """Differentiable cubic-tube coverage with round caps and joins on CUDA."""
+ import torch
+
+ extension = _extension()
+ if (
+ extension is None
+ or not controls.is_cuda
+ or controls.dtype != torch.float32
+ or controls.ndim != 4
+ or controls.shape[1:] != (16, 4, 2)
+ or not widths.is_cuda
+ or widths.dtype != torch.float32
+ or widths.ndim != 1
+ or widths.shape[0] != controls.shape[0]
+ or subpixels not in {1, 2, 4}
+ ):
+ return None
+ left, top, right, bottom = box
+ height, width = bottom - top, right - left
+
+ class StrokeCoverage(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, values, stroke_widths):
+ values = values.contiguous()
+ stroke_widths = stroke_widths.contiguous()
+ ctx.save_for_backward(values, stroke_widths)
+ return extension.stroke_forward(
+ values, stroke_widths, height, width, subpixels, left, top
+ )
+
+ @staticmethod
+ def backward(ctx: Any, *upstreams: Any) -> Any:
+ values, stroke_widths = ctx.saved_tensors
+ upstream = upstreams[0]
+ control_gradients, width_gradients = extension.stroke_backward(
+ values,
+ stroke_widths,
+ upstream.contiguous(),
+ height,
+ width,
+ subpixels,
+ left,
+ top,
+ )
+ return control_gradients, width_gradients
+
+ return StrokeCoverage.apply(controls, widths)
+
+
def multi_coverage_forward(
controls: Any,
offsets: list[int],
diff --git a/src/vectrify/refine/paths.py b/src/vectrify/refine/paths.py
index 10b85c73..4f2d16ec 100644
--- a/src/vectrify/refine/paths.py
+++ b/src/vectrify/refine/paths.py
@@ -17,6 +17,7 @@
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from functools import lru_cache
+from time import perf_counter
from typing import Any
import numpy as np
@@ -234,7 +235,7 @@ def to_path_d(segments) -> str:
def coverage(
control: Any,
- width: float,
+ width: float | Any,
box: tuple[int, int, int, int],
samples: int | None = None,
softness: float = 0.25,
@@ -251,6 +252,21 @@ def coverage(
"""
import torch
+ if control.is_cuda and control.shape[0] <= _FUSED_CUBICS:
+ from vectrify.refine.cuda_renderer import stroke_coverage
+
+ padded = _pad_fused_cubics(control[None])
+ stroke_width = (
+ width.reshape(1)
+ if isinstance(width, torch.Tensor)
+ else torch.full(
+ (1,), float(width), dtype=control.dtype, device=control.device
+ )
+ )
+ native = stroke_coverage(padded, stroke_width, box, subpixels=2)
+ if native is not None:
+ return native[0]
+
if samples is None:
samples = _samples_for(control)
left, top, right, bottom = box
@@ -334,6 +350,7 @@ def fit_group(
widths: float | list[float],
target: Image.Image,
backdrop: Image.Image,
+ colours: list[tuple[float, float, float]] | None = None,
size: int = 700,
steps: int = 200,
samples: int | None = None,
@@ -343,23 +360,20 @@ def fit_group(
redundancy: float = 0.15,
smooth: float = 0.0,
anchor: float = 0.001,
-) -> tuple[list[str], float, float]:
+) -> tuple[list[str], list[float], list[tuple[float, float, float]], float, float]:
"""Fit every path in *paths* together, returning new path data and losses.
*backdrop* is the drawing rendered with this group removed; *target* is the
- picture being matched. Both are greyscale and the same size as the canvas.
+ picture being matched. Both are RGB and the same size as the canvas.
*pinned* names welded vertices that must not move: a point this set shares
with a path outside it. Without them a partial fit tears the drawing at
exactly the junctions welding exists to hold -- the fitted side walks away
while the neighbour it meets stays put.
- The paths composite as a soft union -- one minus the product of their
- complements -- which is what "any of these strokes covers this pixel" means
- and what the real renderer shows. A redundancy term charges for pixels more
- than one stroke covers, because the union alone is indifferent between three
- strokes doing a third of the work each and one doing all of it while the
- other two collapse onto it.
+ Each path is composited in SVG document order over the fixed backdrop. This
+ retains different stroke colours and makes width, colour, and cubic controls
+ jointly differentiable parameters of the same local move.
"""
import torch
@@ -371,7 +385,7 @@ def fit_group(
def crop(image: Image.Image) -> Any:
- array = np.asarray(image.convert("L").resize((size, size)), dtype=np.float32)
+ array = np.asarray(image.convert("RGB").resize((size, size)), dtype=np.float32)
return torch.tensor(array[top:bottom, left:right] / 255.0, device=device)
goal = crop(target)
@@ -385,12 +399,25 @@ def crop(image: Image.Image) -> Any:
if isinstance(widths, int | float)
else list(widths)
)
+ colour_values = colours or [(0.0, 0.0, 0.0)] * len(paths)
+ if len(colour_values) != len(paths):
+ raise ValueError("each stroked path needs one RGB colour")
welded, index = weld(chains)
vertices = torch.tensor(welded, device=device, dtype=torch.float32)
vertices.requires_grad_(True)
+ stroke_widths = torch.tensor(each, device=device, dtype=torch.float32)
+ stroke_widths.requires_grad_(True)
+ stroke_colours = torch.tensor(colour_values, device=device, dtype=torch.float32)
+ stroke_colours.requires_grad_(True)
original = vertices.detach().clone()
rows = [torch.tensor(r, device=device, dtype=torch.long) for r in index]
- optimizer = torch.optim.Adam([vertices], lr=learning_rate)
+ optimizer = torch.optim.Adam(
+ [
+ {"params": [vertices], "lr": learning_rate},
+ {"params": [stroke_widths], "lr": learning_rate * 0.1},
+ {"params": [stroke_colours], "lr": learning_rate * 0.05},
+ ]
+ )
def chain_of(row: Any) -> Any:
return vertices[row]
@@ -402,7 +429,7 @@ def controls_of(chain: Any) -> Any:
mask = _focus_mask(
[
coverage(controls_of(chain_of(r)), w, box, samples=samples)
- for r, w in zip(rows, each, strict=True)
+ for r, w in zip(rows, stroke_widths, strict=True)
],
int(margin),
)
@@ -412,12 +439,15 @@ def controls_of(chain: Any) -> Any:
for step in range(steps):
covers = [
coverage(controls_of(chain_of(r)), w, box, samples=samples)
- for r, w in zip(rows, each, strict=True)
+ for r, w in zip(rows, stroke_widths, strict=True)
]
stacked = torch.stack(covers)
- union = 1 - torch.prod(1 - stacked, dim=0)
- drawn = under * (1 - union)
- loss = ((drawn - goal).abs() * weight).sum()
+ drawn = under
+ for alpha, colour in zip(stacked, stroke_colours, strict=True):
+ drawn = drawn * (1 - alpha[..., None]) + (
+ colour.clamp(0, 1) * alpha[..., None]
+ )
+ loss = ((drawn - goal).abs() * weight[..., None]).sum()
if redundancy:
loss = (
loss + redundancy * ((stacked.sum(0) - 1).clamp_min(0) * weight).sum()
@@ -436,6 +466,9 @@ def controls_of(chain: Any) -> Any:
optimizer.zero_grad()
loss.backward()
optimizer.step()
+ with torch.no_grad():
+ stroke_widths.clamp_(min=0.1)
+ stroke_colours.clamp_(0, 1)
if pinned:
with torch.no_grad():
held = torch.tensor(sorted(pinned), device=device, dtype=torch.long)
@@ -443,7 +476,13 @@ def controls_of(chain: Any) -> Any:
last = float(loss.detach())
fitted = [knots_to_path_d(chain_of(r).detach().cpu().tolist()) for r in rows]
- return fitted, first, last
+ return (
+ fitted,
+ stroke_widths.detach().cpu().tolist(),
+ [tuple(colour) for colour in stroke_colours.detach().cpu().tolist()],
+ first,
+ last,
+ )
def _fill_winding(
@@ -1270,7 +1309,9 @@ def _fill_rgb(value: str | None) -> tuple[float, float, float] | None:
)
-def _composite_opaque_fills(alphas: Any, colours: Any) -> Any:
+def _composite_opaque_fills(
+ alphas: Any, colours: Any, backdrop: Any | None = None
+) -> Any:
"""Composite opaque SVG fills in document order without a layer loop.
Each layer contributes its premultiplied colour through the product of the
@@ -1283,9 +1324,12 @@ def _composite_opaque_fills(alphas: Any, colours: Any) -> Any:
transparency = 1 - alphas
above_inclusive = torch.cumprod(transparency.flip(0), dim=0).flip(0)
above = torch.cat((above_inclusive[1:], torch.ones_like(alphas[:1])), dim=0)
- return (
+ painted = (
colours.clamp(0, 1)[:, None, None, :] * alphas[..., None] * above[..., None]
).sum(dim=0)
+ if backdrop is None:
+ return painted
+ return painted + backdrop * above_inclusive[0][..., None]
@lru_cache(maxsize=1)
@@ -1320,6 +1364,7 @@ def fit_filled_svg(
subpixels: int = 2,
monolithic: bool | None = None,
curve_samples: int | None = None,
+ backdrop: Image.Image | None = None,
) -> str:
"""Optimise opaque filled cubic SVG paths against an RGB target.
@@ -1431,6 +1476,18 @@ def fit_filled_svg(
/ 255.0,
device=device,
)
+ under = (
+ None
+ if backdrop is None
+ else torch.tensor(
+ np.asarray(
+ backdrop.convert("RGB").resize((work_width, work_height)),
+ dtype=np.float32,
+ )
+ / 255.0,
+ device=device,
+ )
+ )
point_optimizer = torch.optim.Adam(
[control_storage], lr=point_learning_rate, fused=device == "cuda"
)
@@ -1797,7 +1854,11 @@ def rasterise_multi_group(
if goal.is_cuda
else _composite_opaque_fills
)
- rendered = composite(alpha_stack, color_storage)
+ rendered = (
+ composite(alpha_stack, color_storage)
+ if under is None
+ else _composite_opaque_fills(alpha_stack, color_storage, under)
+ )
loss = ((rendered - goal) ** 2).mean()
loss = (
loss
@@ -1838,7 +1899,7 @@ def rasterise_multi_group(
initial_alphas[index] = rasterise_multi(index, path)
before: list[Any] = []
- rendered = torch.zeros_like(goal)
+ rendered = torch.zeros_like(goal) if under is None else under
for index, alpha in enumerate(initial_alphas):
assert alpha is not None
colour = color_storage[index]
@@ -1947,6 +2008,355 @@ def layer_loss(
PATH_FIT = "Mutation: path fit"
+def fittable_opaque_fills(svg: str) -> bool:
+ """Whether *svg* contains a fill the analytic CUDA fitter can optimise."""
+ import xml.etree.ElementTree as ET
+
+ try:
+ root = ET.fromstring(svg)
+ except ET.ParseError:
+ return False
+ for element in root.iter():
+ if (
+ element.tag.split("}")[-1] != "path"
+ or _fill_rgb(element.get("fill")) is None
+ ):
+ continue
+ try:
+ parse_filled_cubics(element.get("d", ""))
+ except UnsupportedPathError:
+ continue
+ if element.get("fill-rule", "nonzero").strip().lower() in {
+ "evenodd",
+ "nonzero",
+ }:
+ return True
+ return False
+
+
+_FillBounds = tuple[float, float, float, float]
+_FittableFill = tuple[int, Any, _FillBounds]
+
+
+def _fittable_fill_elements(root) -> list[_FittableFill]:
+ """Return document-indexed opaque fills with conservative control bounds."""
+ entries = []
+ for document_index, element in enumerate(root.iter()):
+ if element.tag.split("}")[-1] != "path" or not element.get("d"):
+ continue
+ if _fill_rgb(element.get("fill")) is None:
+ continue
+ try:
+ contours = parse_filled_cubics(element.get("d", ""))
+ except UnsupportedPathError:
+ continue
+ if element.get("fill-rule", "nonzero").strip().lower() not in {
+ "evenodd",
+ "nonzero",
+ }:
+ continue
+ points = [point for contour in contours for cubic in contour for point in cubic]
+ entries.append(
+ (
+ document_index,
+ element,
+ (
+ min(point[0] for point in points),
+ min(point[1] for point in points),
+ max(point[0] for point in points),
+ max(point[1] for point in points),
+ ),
+ )
+ )
+ return entries
+
+
+def _select_fill_group(
+ entries: list[_FittableFill],
+ *,
+ weights: Mapping[int, float] | None,
+ maximum_paths: int,
+) -> set[int]:
+ """Choose one bounded spatial fill group, biased toward attributed error."""
+ if maximum_paths < 1:
+ raise ValueError("maximum_paths must be positive")
+ scores = [max(0.0, (weights or {}).get(index, 0.0)) for index, _el, _box in entries]
+ focal = (
+ random.choices(entries, weights=scores, k=1)[0]
+ if sum(scores) > 0
+ else random.choice(entries)
+ )
+ focal_index, _element, (left, top, right, bottom) = focal
+ centre_x, centre_y = (left + right) / 2, (top + bottom) / 2
+ extent = max(right - left, bottom - top, 8.0)
+
+ def distance(entry: _FittableFill) -> tuple[int, float, int]:
+ (
+ index,
+ _candidate,
+ (
+ candidate_left,
+ candidate_top,
+ candidate_right,
+ candidate_bottom,
+ ),
+ ) = entry
+ candidate_x = (candidate_left + candidate_right) / 2
+ candidate_y = (candidate_top + candidate_bottom) / 2
+ overlap = not (
+ candidate_right < left - extent
+ or candidate_left > right + extent
+ or candidate_bottom < top - extent
+ or candidate_top > bottom + extent
+ )
+ return (
+ 0 if overlap else 1,
+ (candidate_x - centre_x) ** 2 + (candidate_y - centre_y) ** 2,
+ index,
+ )
+
+ selected = sorted(entries, key=distance)[:maximum_paths]
+ return {index for index, _element, _box in selected} | {focal_index}
+
+
+def fill_groups(svg: str, *, maximum_paths: int = 16) -> list[set[int]]:
+ """Partition opaque fills into bounded spatial groups for coordinate descent."""
+ import xml.etree.ElementTree as ET
+
+ entries = _fittable_fill_elements(ET.fromstring(svg))
+ remaining = {index for index, _element, _box in entries}
+ groups = []
+ while remaining:
+ focal = next(entry for entry in entries if entry[0] in remaining)
+ focal_index, _element, (left, top, right, bottom) = focal
+ centre_x, centre_y = (left + right) / 2, (top + bottom) / 2
+ extent = max(right - left, bottom - top, 8.0)
+
+ def key(
+ entry: _FittableFill,
+ bounds: _FillBounds = (left, top, right, bottom),
+ radius: float = extent,
+ centre: tuple[float, float] = (centre_x, centre_y),
+ ) -> tuple[int, float, int]:
+ (
+ index,
+ _candidate,
+ (
+ candidate_left,
+ candidate_top,
+ candidate_right,
+ candidate_bottom,
+ ),
+ ) = entry
+ focal_left, focal_top, focal_right, focal_bottom = bounds
+ focal_x, focal_y = centre
+ candidate_x = (candidate_left + candidate_right) / 2
+ candidate_y = (candidate_top + candidate_bottom) / 2
+ overlap = not (
+ candidate_right < focal_left - radius
+ or candidate_left > focal_right + radius
+ or candidate_bottom < focal_top - radius
+ or candidate_top > focal_bottom + radius
+ )
+ return (
+ 0 if overlap else 1,
+ (candidate_x - focal_x) ** 2 + (candidate_y - focal_y) ** 2,
+ index,
+ )
+
+ group = {
+ index
+ for index, _element, _bounds in sorted(
+ (entry for entry in entries if entry[0] in remaining), key=key
+ )[:maximum_paths]
+ }
+ group.add(focal_index)
+ groups.append(group)
+ remaining -= group
+ return groups
+
+
+def fit_opaque_fills_locally(
+ svg: str,
+ reference_png: bytes,
+ *,
+ steps: int = 8,
+ rasterize=None,
+ weights: Mapping[int, float] | None = None,
+ maximum_paths: int = 16,
+ selected_indices: set[int] | None = None,
+ optimisation_long_side: int | None = 64,
+ gpu_gate: Any = None,
+) -> str:
+ """Fit one spatially bounded opaque-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
+ optimisation raster used by SAMVG; this is a local move, not its 500-step
+ seed-fitting phase.
+ """
+ from PIL import Image
+
+ target = Image.open(io.BytesIO(reference_png)).convert("RGB")
+ if rasterize is None:
+ raise UnsupportedPathError("bounded fill fitting needs an SVG rasterizer")
+ import xml.etree.ElementTree as ET
+
+ original = ET.fromstring(svg)
+ entries = _fittable_fill_elements(original)
+ if not entries:
+ raise UnsupportedPathError("no opaque filled cubic paths to fit")
+ selected_indices = selected_indices or _select_fill_group(
+ entries, weights=weights, maximum_paths=maximum_paths
+ )
+ backdrop_root = ET.fromstring(svg)
+ working_root = ET.fromstring(svg)
+ for index, element in enumerate(backdrop_root.iter()):
+ if index in selected_indices:
+ element.set("d", "")
+ for index, element in enumerate(working_root.iter()):
+ if index not in selected_indices and element.tag.split("}")[-1] == "path":
+ element.set("d", "")
+ backdrop = Image.open(
+ io.BytesIO(
+ rasterize(
+ ET.tostring(backdrop_root, encoding="unicode"),
+ target.width,
+ target.height,
+ )
+ )
+ ).convert("RGB")
+ with gpu_slot(gpu_gate):
+ fitted = fit_filled_svg(
+ ET.tostring(working_root, encoding="unicode"),
+ target,
+ steps=steps,
+ optimisation_long_side=optimisation_long_side,
+ backdrop=backdrop,
+ )
+ fitted_root = ET.fromstring(fitted)
+ fitted_by_index = dict(enumerate(fitted_root.iter()))
+ for index, element in enumerate(original.iter()):
+ if index not in selected_indices:
+ continue
+ updated = fitted_by_index[index]
+ element.set("d", updated.get("d", ""))
+ element.set("fill", updated.get("fill", element.get("fill", "")))
+ return ET.tostring(original, encoding="unicode")
+
+
+def fit_filled_svg_bounded(
+ svg: str,
+ target: Image.Image,
+ *,
+ rasterize,
+ steps: int = 500,
+ maximum_paths: int = 16,
+ gpu_gate: Any = None,
+ measurements: list[dict[str, int | float]] | None = None,
+) -> str:
+ """Run one full SAMVG fill phase as bounded spatial coordinate descent.
+
+ ``steps`` is the per-group phase budget. 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
+ graph. When requested, ``measurements`` receives one timing and CUDA-peak
+ record for each local group mutation.
+ """
+ if steps < 1:
+ raise ValueError("steps must be positive")
+ groups = fill_groups(svg, maximum_paths=maximum_paths)
+ if not groups:
+ raise UnsupportedPathError("no opaque filled cubic paths to optimise")
+ encoded = io.BytesIO()
+ target.convert("RGB").save(encoded, format="PNG")
+ fitted = svg
+ for index, group in enumerate(groups):
+ 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]
+ started = perf_counter()
+ fitted = fit_opaque_fills_locally(
+ fitted,
+ encoded.getvalue(),
+ steps=steps,
+ rasterize=rasterize,
+ maximum_paths=maximum_paths,
+ selected_indices=group,
+ optimisation_long_side=None,
+ gpu_gate=gpu_gate,
+ )
+ 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": index,
+ "paths": len(group),
+ "seconds": perf_counter() - started,
+ "peak_cuda_bytes": peak,
+ }
+ )
+ return fitted
+
+
+def fittable_strokes(svg: str) -> bool:
+ """Whether the unified cubic-stroke fitter can select a stroke group."""
+ import xml.etree.ElementTree as ET
+
+ try:
+ return bool(fittable_clusters(ET.fromstring(svg)))
+ except ET.ParseError:
+ return False
+
+
+def fit_svg_primitives_locally(
+ svg: str,
+ reference_png: bytes,
+ *,
+ rasterize,
+ weights: Mapping[int, float] | None = None,
+ steps: int = 8,
+ gpu_gate: Any = None,
+) -> str:
+ """Fit one selected fill or stroke primitive group over fixed SVG context.
+
+ Each fitter rasterizes the non-active document as a fixed backdrop. This
+ prevents a fill from being rewarded for covering a line or editable text,
+ while the subsequent stroke move sees the newly fitted fills unchanged.
+ """
+ fills = fittable_opaque_fills(svg)
+ strokes = fittable_strokes(svg)
+ if fills and (not strokes or random.random() < 0.5):
+ return fit_opaque_fills_locally(
+ svg,
+ reference_png,
+ steps=steps,
+ rasterize=rasterize,
+ weights=weights,
+ gpu_gate=gpu_gate,
+ )
+ if strokes:
+ return fit_random_group(
+ svg,
+ reference_png,
+ rasterize=rasterize,
+ steps=steps,
+ weights=weights,
+ gpu_gate=gpu_gate,
+ )
+ raise UnsupportedPathError("no supported filled or stroked cubics to fit")
+
+
def _stroke_width(element, ancestors) -> float | None:
"""Return the inherited stroke width, or ``None`` for an unpainted path."""
for node in (element, *ancestors):
@@ -1961,6 +2371,20 @@ def _stroke_width(element, ancestors) -> float | None:
return None
+def _stroke_rgb(
+ element: Any, ancestors: list[Any]
+) -> tuple[float, float, float] | None:
+ """Return an inherited opaque hex stroke colour, if the path paints one."""
+ for node in (element, *ancestors):
+ raw = node.get("stroke")
+ if raw is None:
+ continue
+ if raw.strip().lower() == "none":
+ return None
+ return _fill_rgb(raw)
+ return None
+
+
def _parents(root) -> dict[int, Any]:
"""id(child) -> parent, so a path can be read in the context it inherits."""
table: dict[int, Any] = {}
@@ -2084,7 +2508,7 @@ def fit_random_group(
paths = [p for i, p in enumerate(paths) if i in chosen]
size = int(_canvas_side(root))
- target = Image.open(io.BytesIO(reference_png)).convert("L").resize((size, size))
+ target = Image.open(io.BytesIO(reference_png)).convert("RGB").resize((size, size))
# The backdrop is the drawing without these paths, so the fit sees the rest
# of the picture as a constant and cannot be rewarded for redrawing it.
@@ -2095,24 +2519,37 @@ def fit_random_group(
path.set("d", "")
backdrop = Image.open(
io.BytesIO(rasterize(ET.tostring(root, encoding="unicode"), size, size))
- ).convert("L")
+ ).convert("RGB")
for path, data in zip(paths, original, strict=True):
path.set("d", data)
held = _shared_vertices([parse_cubics(d) for d in original], excluded)
+ parents = _parents(root)
+ colours = [
+ _stroke_rgb(path, _ancestry(path, parents, root)) or (0.0, 0.0, 0.0)
+ for path in paths
+ ]
with gpu_slot(gpu_gate):
- fitted, _first, _last = fit_group(
+ fitted, fitted_widths, fitted_colours, _first, _last = fit_group(
original,
widths,
target,
backdrop,
+ colours,
size=size,
steps=steps,
samples=samples,
pinned=held,
)
- for path, data in zip(paths, fitted, strict=True):
+ for path, data, width, colour in zip(
+ paths, fitted, fitted_widths, fitted_colours, strict=True
+ ):
path.set("d", data)
+ path.set("stroke-width", f"{width:.2f}")
+ path.set(
+ "stroke",
+ "#" + "".join(f"{round(channel * 255):02x}" for channel in colour),
+ )
return ET.tostring(root, encoding="unicode")
diff --git a/src/vectrify/refine/samvg.py b/src/vectrify/refine/samvg.py
index 115ed43b..fa0b8b39 100644
--- a/src/vectrify/refine/samvg.py
+++ b/src/vectrify/refine/samvg.py
@@ -9,13 +9,18 @@
from __future__ import annotations
import io
+import itertools
+import json
import logging
import math
import os
+import re
import xml.etree.ElementTree as ET
from collections import defaultdict
+from collections.abc import Callable
+from contextlib import nullcontext
from dataclasses import dataclass
-from typing import cast
+from typing import Any, cast
import numpy as np
from PIL import Image
@@ -26,11 +31,30 @@
# ViT-H is the paper-quality default; users who need the smaller checkpoint can
# opt down without changing the package through VECTRIFY_SAMVG_MODEL.
SAMVG_MODEL = os.environ.get("VECTRIFY_SAMVG_MODEL", "facebook/sam-vit-huge")
+# SAM encodes images at a native 1024px long side. Keep that encoder-size cap
+# as the default even when Vectrify is asked to vectorize a larger original;
+# masks are restored to the original canvas before tracing.
+SAMVG_MAX_SIDE = int(os.environ.get("VECTRIFY_SAMVG_MAX_SIDE", "1024"))
+# This is the decoder prompt batch, not the dissertation's 32x32 sampling
+# 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
+# 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
+# them in the source coordinate system.
+SAMVG_OCR_MODEL = os.environ.get(
+ "VECTRIFY_SAMVG_OCR_MODEL", "Qwen/Qwen2.5-VL-3B-Instruct"
+)
+# OCR text is often a few pixels off because its original font is unknown.
+# Permit that small mismatch (per affected channel), but never a large visual
+# regression just because the VLM claimed confidence.
+OCR_TEXT_RMSE_TOLERANCE = 0.02
@dataclass(frozen=True)
@@ -43,6 +67,166 @@ class MaskLayer:
overlap_pixels: int = 0
+@dataclass(frozen=True)
+class TextLayer:
+ """A high-confidence OCR word represented as editable SVG text."""
+
+ text: str
+ x: float
+ y: float
+ width: float
+ height: float
+ colour: tuple[int, int, int]
+ angle: float = 0.0
+
+
+def _text_colour(pixels: np.ndarray) -> tuple[int, int, int]:
+ """Estimate ink colour by contrasting a word crop with its border."""
+ height, width, _channels = pixels.shape
+ if height < 3 or width < 3:
+ colour = pixels.reshape(-1, 3).mean(axis=0)
+ else:
+ border = np.concatenate(
+ (pixels[0], pixels[-1], pixels[1:-1, 0], pixels[1:-1, -1])
+ )
+ background = border.mean(axis=0)
+ distance = np.linalg.norm(pixels.astype(np.float32) - background, axis=2)
+ ink = pixels[distance >= np.percentile(distance, 80)]
+ colour = ink.mean(axis=0) if len(ink) else background
+ return cast(tuple[int, int, int], tuple(int(value) for value in np.rint(colour)))
+
+
+def _ocr_json(response: str) -> list[dict[str, object]]:
+ """Decode the strict JSON array requested from the vision-language model."""
+ match = re.search(r"\[[\s\S]*\]", response)
+ if match is None:
+ return []
+ try:
+ parsed = json.loads(match.group())
+ except json.JSONDecodeError:
+ return []
+ if not isinstance(parsed, list):
+ return []
+ return [item for item in parsed if isinstance(item, dict)]
+
+
+def detect_text(image: Image.Image, *, confidence: float = 0.8) -> list[TextLayer]:
+ """Read editable text using Qwen2.5-VL's 3B Torch model.
+
+ It returns content and source-pixel bounding boxes in one inference pass.
+ We keep only the VLM's high-confidence multi-character labels: a guessed
+ font is worse than the normal SAMVG filled-path representation.
+ """
+ try:
+ import torch
+ from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+ except ImportError as exc: # pragma: no cover - installation-specific
+ raise ImportError(
+ "SAMVG OCR requires the samvg extra. Install 'vectrify[samvg]'."
+ ) from exc
+ source = np.asarray(image.convert("RGB"))
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ dtype = torch.bfloat16 if device == "cuda" else torch.float32
+ prompt = (
+ "Read visible text in this image. Return only a JSON array. Each entry "
+ 'must be {"text": string, "box": [left, top, right, bottom], '
+ '"confidence": number}. Boxes must use this image\'s pixel '
+ "coordinates. Include only clearly readable labels of at least two "
+ "characters, and do not describe icons, logos, or non-text shapes."
+ )
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "image": image},
+ {"type": "text", "text": prompt},
+ ],
+ }
+ ]
+ processor = AutoProcessor.from_pretrained(SAMVG_OCR_MODEL)
+ # Transformers currently exposes a descriptor mismatch between this model
+ # class and GenerationMixin to Pyrefly; runtime generation is the normal
+ # PreTrainedModel API.
+ model: Any = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+ SAMVG_OCR_MODEL, torch_dtype=dtype
+ ).to(device)
+ detected: list[TextLayer] = []
+ try:
+ chat = processor.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+ )
+ inputs = processor(
+ text=[chat], images=[image], padding=True, return_tensors="pt"
+ ).to(device)
+ with torch.inference_mode():
+ output = model.generate(**inputs, max_new_tokens=768, do_sample=False)
+ generated = output[:, inputs.input_ids.shape[1] :]
+ response = processor.batch_decode(
+ generated, skip_special_tokens=True, clean_up_tokenization_spaces=False
+ )[0]
+ for entry in _ocr_json(response):
+ text = entry.get("text")
+ box = entry.get("box")
+ score = entry.get("confidence")
+ if (
+ not isinstance(text, str)
+ or not isinstance(box, list)
+ or len(box) != 4
+ or not isinstance(score, (int, float))
+ or float(score) < confidence
+ or len(text.strip()) < 2
+ ):
+ continue
+ try:
+ x, y, right, bottom = (float(value) for value in box)
+ except (TypeError, ValueError):
+ continue
+ x, y = max(0.0, x), max(0.0, y)
+ right = min(float(image.width), right)
+ bottom = min(float(image.height), bottom)
+ width, height = right - x, bottom - y
+ if width < 4 or height < 4:
+ continue
+ crop = source[
+ math.floor(y) : math.ceil(bottom), math.floor(x) : math.ceil(right)
+ ]
+ if not crop.size:
+ continue
+ detected.append(
+ TextLayer(
+ text=text.strip(),
+ x=x,
+ y=y,
+ width=width,
+ height=height,
+ colour=_text_colour(crop),
+ )
+ )
+ finally:
+ del model
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ log.info("SAMVG OCR: retained %d editable text layer(s).", len(detected))
+ return detected
+
+
+def _text_svg_attributes(layer: TextLayer) -> dict[str, str]:
+ """Map OCR geometry to a portable editable SVG text element."""
+ colour = f"#{layer.colour[0]:02x}{layer.colour[1]:02x}{layer.colour[2]:02x}"
+ attributes = {
+ "x": f"{layer.x:.2f}",
+ "y": f"{layer.y + layer.height * 0.8:.2f}",
+ "font-family": "sans-serif",
+ "font-size": f"{layer.height:.2f}",
+ "fill": colour,
+ }
+ if abs(layer.angle) > 1:
+ attributes["transform"] = (
+ f"rotate({layer.angle:.2f} {layer.x:.2f} {layer.y:.2f})"
+ )
+ return attributes
+
+
def _is_crop_edge_mask(
mask: np.ndarray,
crop_box: tuple[int, int, int, int],
@@ -66,55 +250,389 @@ def _is_crop_edge_mask(
return bool(np.any(at_crop_edge & ~at_image_edge))
-def automatic_masks(image: Image.Image) -> list[np.ndarray]:
- """Retrieve SAM AMG masks with the thesis's 32-point grid and crops."""
+def _run_components(mask: np.ndarray) -> list[list[tuple[int, int, int]]]:
+ """Return 4-connected components as row spans in row-major order.
+
+ The old breadth-first walk crossed the Python interpreter once for every
+ foreground pixel. SAM masks are usually broad regions, so representing
+ each row as contiguous runs reduces that to a small number of intervals
+ while retaining scipy.ndimage's 4-connected ordering.
+ """
+ foreground = np.asarray(mask, dtype=bool)
+ _height, width = foreground.shape
+ parent = [0]
+
+ def root(index: int) -> int:
+ while parent[index] != index:
+ parent[index] = parent[parent[index]]
+ index = parent[index]
+ return index
+
+ def merge(left: int, right: int) -> None:
+ left, right = root(left), root(right)
+ if left != right:
+ parent[right] = left
+
+ rows: list[list[tuple[int, int, int]]] = []
+ previous: list[tuple[int, int, int]] = []
+ for row in foreground:
+ padded = np.empty(width + 2, dtype=bool)
+ padded[0] = padded[-1] = False
+ padded[1:-1] = row
+ edges = np.flatnonzero(padded[1:] != padded[:-1])
+ current: list[tuple[int, int, int]] = []
+ prior = 0
+ for start, end in edges.reshape(-1, 2):
+ while prior < len(previous) and previous[prior][1] <= start:
+ prior += 1
+ index = len(parent)
+ parent.append(index)
+ candidate = prior
+ while candidate < len(previous) and previous[candidate][0] < end:
+ merge(index, previous[candidate][2])
+ candidate += 1
+ current.append((int(start), int(end), index))
+ rows.append(current)
+ previous = current
+
+ components: list[list[tuple[int, int, int]]] = []
+ component_ids: dict[int, int] = {}
+ for y, runs in enumerate(rows):
+ for start, end, index in runs:
+ component = root(index)
+ label = component_ids.setdefault(component, len(component_ids))
+ if label == len(components):
+ components.append([])
+ components[label].append((y, start, end))
+ return components
+
+
+def _label(mask: np.ndarray) -> tuple[np.ndarray, int]:
+ """Materialize 4-connected scanline components as an integer label map."""
+ foreground = np.asarray(mask, dtype=bool)
+ labels = np.zeros(foreground.shape, dtype=np.int32)
+ components = _run_components(foreground)
+ for index, runs in enumerate(components, start=1):
+ for y, start, end in runs:
+ labels[y, start:end] = index
+ return labels, len(components)
+
+
+def _edt_1d(values: np.ndarray) -> np.ndarray:
+ """Squared lower envelope for the linear-time Euclidean distance transform."""
+ size = len(values)
+ infinity = np.inf
+ sites = np.flatnonzero(np.isfinite(values))
+ if not len(sites):
+ return np.full(size, infinity, dtype=np.float64)
+ vertices = np.empty(len(sites), dtype=np.int32)
+ intersections = np.empty(len(sites) + 1, dtype=np.float64)
+ count = 0
+ vertices[0] = sites[0]
+ intersections[0], intersections[1] = -infinity, infinity
+ for site in sites[1:]:
+ intersection = (
+ (values[site] + site * site)
+ - (values[vertices[count]] + vertices[count] * vertices[count])
+ ) / (2 * (site - vertices[count]))
+ while intersection <= intersections[count]:
+ count -= 1
+ intersection = (
+ (values[site] + site * site)
+ - (values[vertices[count]] + vertices[count] * vertices[count])
+ ) / (2 * (site - vertices[count]))
+ count += 1
+ vertices[count] = site
+ intersections[count], intersections[count + 1] = intersection, infinity
+ output = np.empty(size, dtype=np.float64)
+ index = 0
+ for position in range(size):
+ while intersections[index + 1] < position:
+ index += 1
+ site = vertices[index]
+ output[position] = (position - site) ** 2 + values[site]
+ return output
+
+
+def _distance_transform_edt(mask: np.ndarray) -> np.ndarray:
+ """Exact CPU Euclidean distance to the nearest false pixel, without SciPy."""
+ foreground = np.asarray(mask, dtype=bool)
+ height, width = foreground.shape
+ squared = np.where(foreground, np.inf, 0.0)
+ if not np.isfinite(squared).any():
+ yy, xx = np.indices((height, width), dtype=np.float64)
+ return np.hypot(yy + 1, xx)
+ columns = np.empty_like(squared)
+ for column in range(width):
+ columns[:, column] = _edt_1d(squared[:, column])
+ output = np.empty_like(squared)
+ for row in range(height):
+ output[row] = _edt_1d(columns[row])
+ return np.sqrt(output)
+
+
+def _binary_dilation(mask: np.ndarray, iterations: int) -> np.ndarray:
+ """Apply scipy's default 4-connected binary dilation with Torch kernels."""
+ if iterations <= 0:
+ return np.asarray(mask, dtype=bool)
+ import torch
+ import torch.nn.functional as functional
+
+ source = torch.as_tensor(mask, dtype=torch.float32)[None, None]
+ cross = source.new_tensor([[[[0, 1, 0], [1, 1, 1], [0, 1, 0]]]])
+ for _ in range(iterations):
+ source = (functional.conv2d(source, cross, padding=1) > 0).to(source.dtype)
+ return source[0, 0].bool().numpy()
+
+
+def _mean_shift_centres(points: np.ndarray, bandwidth: float) -> np.ndarray:
+ """Deterministic bin-seeded mean shift matching SAMVG's prompt clustering."""
+ bins = np.unique(np.rint(points / bandwidth).astype(np.int32), axis=0)
+ seeds = bins.astype(np.float64) * bandwidth
+ centres: dict[tuple[float, float], int] = {}
+ for seed in seeds:
+ centre = seed
+ members = np.empty(0, dtype=np.int64)
+ for _ in range(300):
+ delta = points - centre
+ members = np.flatnonzero((delta * delta).sum(axis=1) <= bandwidth**2)
+ if not len(members):
+ break
+ updated = points[members].mean(axis=0)
+ if np.linalg.norm(updated - centre) < bandwidth * 1e-3:
+ centre = updated
+ break
+ centre = updated
+ if len(members):
+ centres[tuple(centre)] = len(members)
+ # This intentionally follows sklearn's intensity-then-coordinate ordering
+ # and radius duplicate suppression, preserving the old prompt priority.
+ ordered = sorted(centres.items(), key=lambda item: (item[1], item[0]), reverse=True)
+ candidates = np.asarray([centre for centre, _count in ordered], dtype=np.float64)
+ unique = np.ones(len(candidates), dtype=bool)
+ for index, centre in enumerate(candidates):
+ if unique[index]:
+ neighbours = np.linalg.norm(candidates - centre, axis=1) <= bandwidth
+ unique[neighbours] = False
+ unique[index] = True
+ return candidates[unique]
+
+
+def _sam_image(image: Image.Image, max_side: int | None) -> tuple[Image.Image, float]:
+ """Bound a SAM pass while retaining masks in the original canvas space."""
+ image = image.convert("RGB")
+ if max_side is None:
+ return image, 1.0
+ if max_side < 1:
+ raise ValueError("max_side must be positive")
+ longest = max(image.size)
+ if longest <= max_side:
+ return image, 1.0
+ scale = max_side / longest
+ return (
+ image.resize(
+ (round(image.width * scale), round(image.height * scale)),
+ Image.Resampling.LANCZOS,
+ ),
+ scale,
+ )
+
+
+def _restore_mask(mask: np.ndarray, size: tuple[int, int]) -> np.ndarray:
+ """Nearest-neighbour restore keeps SAM's binary mask semantics."""
+ if mask.shape == (size[1], size[0]):
+ return np.asarray(mask, dtype=bool)
+ return np.asarray(
+ Image.fromarray(np.asarray(mask, dtype=np.uint8) * 255).resize(
+ size, Image.Resampling.NEAREST
+ ),
+ dtype=bool,
+ )
+
+
+@dataclass
+class _SamRuntime:
+ """One SAM model lifetime, including a reusable full-image embedding."""
+
+ generator: Any
+ processor: Any | None = None
+ image_embeddings: Any | None = None
+ embedding_size: tuple[int, int] | None = None
+
+
+def _sam_runtime() -> _SamRuntime:
+ """Load SAM once, in half precision when CUDA is available."""
try:
+ import torch
from transformers import pipeline
except ImportError as exc: # pragma: no cover - installation-specific
raise ImportError(
- "SAMVG requires the vision extra. Install 'vectrify[vision]'."
+ "SAMVG requires the samvg extra. Install 'vectrify[samvg]'."
) from exc
- image = image.convert("RGB")
- generator = pipeline("mask-generation", model=SAMVG_MODEL, device=0)
- log.info("SAMVG automatic masks: %s on %s.", SAMVG_MODEL, generator.device)
+ options: dict[str, Any] = {"model": SAMVG_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,
+ generator.device,
+ "fp16" if torch.cuda.is_available() else "fp32",
+ )
+ return _SamRuntime(generator)
- def masks_for(source: Image.Image) -> list[np.ndarray]:
- return [
- np.asarray(mask, dtype=bool)
- for mask in generator(
- source,
- points_per_batch=32,
- points_per_crop=32,
- crops_n_layers=0,
- pred_iou_thresh=SAMVG_PRED_IOU_THRESH,
- stability_score_thresh=SAMVG_STABILITY_SCORE_THRESH,
- )["masks"]
- ]
+
+def _sam_autocast():
+ """Use Tensor Cores for inference while keeping exported masks binary."""
+ import torch
+
+ if torch.cuda.is_available():
+ return torch.autocast(device_type="cuda", dtype=torch.float16)
+ return nullcontext()
+
+
+def _automatic_forward(inputs: Any, runtime: _SamRuntime) -> dict[str, Any]:
+ """Decode on CUDA, then expand and filter masks on CPU.
+
+ 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()
+ 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():
+ 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,
+ )
+ return {
+ "masks": filtered_masks,
+ "is_last": is_last,
+ "boxes": boxes,
+ "iou_scores": scores,
+ }
+
+
+def _automatic_masks_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.
+
+ 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.
+ """
+ generator = runtime.generator
+ arguments = {
+ "points_per_batch": points_per_batch,
+ "points_per_crop": 32,
+ "crops_n_layers": 0,
+ "pred_iou_thresh": SAMVG_PRED_IOU_THRESH,
+ "stability_score_thresh": SAMVG_STABILITY_SCORE_THRESH,
+ }
+ # 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"]]
+
+ outputs = []
+ for inputs in generator.preprocess(
+ source,
+ points_per_batch=points_per_batch,
+ points_per_crop=32,
+ crops_n_layers=0,
+ ):
+ # ChunkPipeline normally performs this transfer between preprocess and
+ # _forward. We call those stages directly to retain the embedding.
+ inputs = generator._ensure_tensor_on_device(inputs, device=generator.device)
+ embedding = inputs.get("image_embeddings")
+ if (
+ cache_embedding
+ and embedding is not None
+ and runtime.image_embeddings is None
+ ):
+ 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"]]
+
+
+def automatic_masks(
+ image: Image.Image,
+ *,
+ max_side: int | None = SAMVG_MAX_SIDE,
+ _runtime: _SamRuntime | None = None,
+) -> list[np.ndarray]:
+ """Retrieve SAM AMG masks with the thesis grid, optionally size-capped."""
+ original_size = image.size
+ image, _scale = _sam_image(image, max_side)
+ runtime = _runtime or _sam_runtime()
# transformers' built-in crop layer tries to stack unequal crop tensors.
# Run that first crop layer one crop at a time instead. Crucially, do not
# pre-pad a rectangular image: the original AMG formula uses the source's
# short side for overlap, and black padding changes SAM's visual context.
width, height = image.size
- collected = masks_for(image)
- 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 {
- (0, 0),
- (crop_width - overlap, 0),
- (0, crop_height - overlap),
- (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 masks_for(image.crop(crop_box)):
- if _is_crop_edge_mask(crop_mask, crop_box, image.size):
- continue
- mask = np.zeros((height, width), dtype=bool)
- mask[y:bottom, x:right] = crop_mask
- collected.append(mask)
- return collected
+
+ def collect(points_per_batch: int) -> list[np.ndarray]:
+ collected = _automatic_masks_for(
+ image,
+ runtime,
+ cache_embedding=True,
+ points_per_batch=points_per_batch,
+ )
+ 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 {
+ (0, 0),
+ (crop_width - overlap, 0),
+ (0, crop_height - overlap),
+ (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(
+ 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):
+ continue
+ mask = np.zeros((height, width), dtype=bool)
+ mask[y:bottom, x:right] = crop_mask
+ collected.append(mask)
+ return collected
+
+ collected = collect(SAMVG_POINTS_PER_BATCH)
+ return [_restore_mask(mask, original_size) for mask in collected]
def _components(
@@ -126,30 +644,33 @@ def _components(
holes before tracing matches AMG's small-region cleanup and prevents a
noisy mask from becoming hundreds of even-odd SVG contours.
"""
- from scipy.ndimage import label
-
- labels, count = label(mask)
+ foreground = np.asarray(mask, dtype=bool)
+ if int(foreground.sum()) < min_pixels:
+ return []
+ height, width = foreground.shape
components = []
- for index in range(1, count + 1):
- component = labels == index
- if int(component.sum()) < min_pixels:
+ for runs in _run_components(foreground):
+ if sum(end - start for _y, start, end in runs) < min_pixels:
continue
+ component = np.zeros((height, width), dtype=bool)
+ for y, start, end in runs:
+ component[y, start:end] = True
if fill_holes:
# 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.
- background, hole_count = label(~component)
- for hole in range(1, hole_count + 1):
- points = background == hole
- if (
- int(points.sum()) <= min_pixels
- and not points[0].any()
- and not points[-1].any()
- and not points[:, 0].any()
- and not points[:, -1].any()
- ):
- component[points] = True
+ for hole in _run_components(~component):
+ 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
+ for y, start, end in hole
+ )
+ if not touches_border:
+ for y, start, end in hole:
+ component[y, start:end] = True
components.append(np.asarray(component, dtype=bool))
return components
@@ -195,13 +716,20 @@ def recolour_visible_layers(
return list(reversed(revised))
-def _impact_error(
+def _impact_error_map(
target: np.ndarray, canvas: np.ndarray, coverage: np.ndarray
-) -> float:
+) -> np.ndarray:
"""SAMVG's blank-canvas error, charging uncovered pixels maximally."""
error = ((target.astype(np.float32) - canvas.astype(np.float32)) / 255.0) ** 2
error[~coverage] = 1.0
- return float(error.mean())
+ return error
+
+
+def _impact_error(
+ target: np.ndarray, canvas: np.ndarray, coverage: np.ndarray
+) -> float:
+ """Return the scalar blank-canvas reconstruction error."""
+ return float(_impact_error_map(target, canvas, coverage).mean())
def filter_by_impact(
@@ -234,7 +762,9 @@ def filter_by_impact(
if initial_coverage.shape != coverage.shape:
raise ValueError("initial coverage does not match the target size")
coverage = initial_coverage.astype(bool, copy=True)
- error = _impact_error(target, canvas, coverage)
+ 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
@@ -252,15 +782,22 @@ def filter_by_impact(
tuple[int, int, int],
tuple(int(value) for value in np.rint(target[mask].mean(axis=0))),
)
- next_canvas = canvas.copy()
- next_coverage = coverage | mask
- next_canvas[mask] = colour
- next_error = _impact_error(target, next_canvas, next_coverage)
+ old_error = error_map[mask]
+ next_error_values = (
+ (target[mask].astype(np.float32) - np.asarray(colour, dtype=np.float32))
+ / 255.0
+ ) ** 2
+ next_error_total = error_total - float(old_error.sum(dtype=np.float64))
+ next_error_total += float(next_error_values.sum(dtype=np.float64))
+ next_error = next_error_total / error_map.size
impact = error - next_error
if impact < min_impact:
continue
accepted.append(MaskLayer(mask, colour, impact))
- canvas, coverage, error = next_canvas, next_coverage, next_error
+ canvas[mask] = colour
+ coverage |= mask
+ error_map[mask] = next_error_values
+ error_total, error = next_error_total, next_error
# 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.
@@ -277,18 +814,15 @@ def coverage_prompt_points(
max_points: int = 16,
) -> list[tuple[int, int]]:
"""Find mean-shift centres of large circles untouched by retained masks."""
- from scipy.ndimage import distance_transform_edt
- from sklearn.cluster import MeanShift
-
_canvas, coverage = _render_layers(shape, layers)
radius = max(2, round(min(shape) * radius_fraction))
- distance = np.asarray(distance_transform_edt(~coverage))
+ distance = _distance_transform_edt(~coverage)
ys, xs = np.nonzero(distance >= radius)
if len(xs) == 0:
return []
stride = max(1, len(xs) // 2_048)
points = np.column_stack((xs[::stride], ys[::stride]))
- centres = MeanShift(bandwidth=radius, bin_seeding=True).fit(points).cluster_centers_
+ centres = _mean_shift_centres(points, radius)
ranked = sorted(
((float(distance[round(y), round(x)]), round(x), round(y)) for x, y in centres),
reverse=True,
@@ -297,7 +831,11 @@ def coverage_prompt_points(
def prompted_masks(
- image: Image.Image, points: list[tuple[int, int]]
+ image: Image.Image,
+ points: list[tuple[int, int]],
+ *,
+ max_side: int | None = SAMVG_MAX_SIDE,
+ _runtime: _SamRuntime | None = None,
) -> list[np.ndarray]:
"""Prompt SAM at centres and return all three masks per point.
@@ -308,32 +846,47 @@ def prompted_masks(
if not points:
return []
import torch
- from transformers import SamModel, SamProcessor
-
+ from transformers import SamProcessor
+
+ original_size = image.size
+ image, scale = _sam_image(image, max_side)
+ if scale != 1.0:
+ points = [(round(x * scale), round(y * scale)) for x, y in points]
+ own_runtime = _runtime is None
+ runtime = _runtime or _sam_runtime()
device = "cuda" if torch.cuda.is_available() else "cpu"
log.info("SAMVG prompted masks: using %s.", device)
- processor = SamProcessor.from_pretrained(SAMVG_MODEL)
- model = SamModel.from_pretrained(SAMVG_MODEL).to(device)
+ if runtime.processor is None:
+ runtime.processor = SamProcessor(runtime.generator.image_processor)
try:
input_points = [[[list(point)] for point in points]]
- inputs = processor(
+ inputs = runtime.processor(
images=image, input_points=input_points, return_tensors="pt"
).to(device)
- with torch.inference_mode():
- output = model(**inputs)
- post = processor.image_processor.post_process_masks(
+ 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 [
- np.asarray(post[prompt, candidate], dtype=bool)
+ _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])
]
finally:
- del model
- if torch.cuda.is_available():
+ if own_runtime and torch.cuda.is_available():
torch.cuda.empty_cache()
@@ -345,10 +898,17 @@ def retrieve_layers(
min_impact: float = 1e-5,
max_layers: int = 512,
fill_holes: bool = True,
+ max_side: int | None = SAMVG_MAX_SIDE,
+ _runtime: _SamRuntime | None = None,
) -> list[MaskLayer]:
"""Run SAMVG's automatic-mask, coverage-prompt, filter sequence."""
image = image.convert("RGB")
- initial = automatic_masks(image) if masks is None else masks
+ runtime = _runtime
+ if masks is None:
+ runtime = runtime or _sam_runtime()
+ initial = automatic_masks(image, max_side=max_side, _runtime=runtime)
+ else:
+ initial = masks
layers = filter_by_impact(
image,
initial,
@@ -359,7 +919,7 @@ def retrieve_layers(
)
layers = recolour_visible_layers(image, layers)
points = coverage_prompt_points(layers, (image.height, image.width))
- prompted = prompted_masks(image, points)
+ prompted = prompted_masks(image, points, max_side=max_side, _runtime=runtime)
recovered = filter_by_impact(
image,
prompted,
@@ -541,13 +1101,284 @@ def mask_path(
) -> str | None:
"""Fit every mask contour as a fixed-count cubic Bezier SVG path."""
if overlap_pixels:
- from scipy.ndimage import binary_dilation
-
- mask = binary_dilation(mask, iterations=overlap_pixels)
+ mask = _binary_dilation(mask, overlap_pixels)
parts = [piece for loop in _loops(mask) if (piece := _cubic_loop(loop, segments))]
return " ".join(parts) or None
+_SKELETON_NEIGHBOURS = (
+ (-1, -1),
+ (-1, 0),
+ (-1, 1),
+ (0, -1),
+ (0, 1),
+ (1, -1),
+ (1, 0),
+ (1, 1),
+)
+
+
+def _thin_mask(mask: np.ndarray) -> np.ndarray:
+ """Zhang--Suen thinning without adding a SciPy/skimage dependency."""
+ thin = np.pad(mask.astype(np.uint8), 1).copy()
+ changed = True
+ while changed:
+ changed = False
+ for phase in range(2):
+ remove: list[tuple[int, int]] = []
+ for y, x in zip(*np.nonzero(thin), strict=True):
+ if y in {0, thin.shape[0] - 1} or x in {0, thin.shape[1] - 1}:
+ continue
+ ring = [
+ thin[y - 1, x],
+ thin[y - 1, x + 1],
+ thin[y, x + 1],
+ thin[y + 1, x + 1],
+ thin[y + 1, x],
+ thin[y + 1, x - 1],
+ thin[y, x - 1],
+ thin[y - 1, x - 1],
+ ]
+ count = sum(ring)
+ transitions = sum(
+ left == 0 and right == 1
+ for left, right in zip(ring, [*ring[1:], ring[0]], strict=True)
+ )
+ if not (2 <= count <= 6 and transitions == 1):
+ continue
+ north, east, south, west = ring[0], ring[2], ring[4], ring[6]
+ blocked = (
+ (north and east and south) or (east and south and west)
+ if phase == 0
+ else (north and east and west) or (north and south and west)
+ )
+ if not blocked:
+ remove.append((y, x))
+ if remove:
+ changed = True
+ for y, x in remove:
+ thin[y, x] = 0
+ return thin[1:-1, 1:-1].astype(bool)
+
+
+def _skeleton_traces(mask: np.ndarray) -> list[np.ndarray]:
+ """Split a thinned medial-axis graph into its endpoint/junction traces."""
+ points = {tuple(point) for point in np.argwhere(_thin_mask(mask))}
+ if len(points) < 2:
+ return []
+
+ def adjacent(point: tuple[int, int]) -> list[tuple[int, int]]:
+ y, x = point
+ output = []
+ for dy, dx in _SKELETON_NEIGHBOURS:
+ candidate = y + dy, x + dx
+ if candidate not in points:
+ continue
+ # A diagonal across an orthogonal staircase is not another graph
+ # edge. Keeping it creates artificial triangles and turns every
+ # curved pixel line into a forest of tiny branches.
+ if dy and dx and ((y + dy, x) in points or (y, x + dx) in points):
+ continue
+ output.append(candidate)
+ return output
+
+ nodes = {point for point in points if len(adjacent(point)) != 2}
+ # Closed loops are better represented by SAMVG's filled path: an open
+ # stroke would introduce caps and a stroke-only loop has no stable start.
+ if not nodes:
+ return []
+ traversed: set[tuple[tuple[int, int], tuple[int, int]]] = set()
+
+ def edge_key(
+ first: tuple[int, int], second: tuple[int, int]
+ ) -> tuple[tuple[int, int], tuple[int, int]]:
+ return (first, second) if first <= second else (second, first)
+
+ traces: list[np.ndarray] = []
+ for node in nodes:
+ for neighbour in adjacent(node):
+ edge = edge_key(node, neighbour)
+ if edge in traversed:
+ continue
+ trace, previous, current = [node], node, neighbour
+ traversed.add(edge)
+ while current not in nodes:
+ trace.append(current)
+ choices = [point for point in adjacent(current) if point != previous]
+ if len(choices) != 1:
+ trace = []
+ break
+ previous, current = current, choices[0]
+ traversed.add(edge_key(previous, current))
+ if trace:
+ trace.append(current)
+ if len(trace) >= 2:
+ traces.append(
+ np.asarray([(x, y) for y, x in trace], dtype=np.float64)
+ )
+ return traces
+
+
+def _trace_path_data(trace: np.ndarray, segments: int) -> str:
+ """Fit multiple cubic sections to a skeleton rather than one global PCA line."""
+ count = max(1, min(segments, math.ceil((len(trace) - 1) / 8)))
+ boundaries = np.linspace(0, len(trace) - 1, count + 1, dtype=int)
+ output = [f"M {trace[0, 0]:.2f} {trace[0, 1]:.2f}"]
+ for first, last in itertools.pairwise(boundaries):
+ sample = trace[first : last + 1]
+ if len(sample) == 2:
+ output.append(f"L {sample[-1, 0]:.2f} {sample[-1, 1]:.2f}")
+ else:
+ control_a, control_b = _fit_cubic(sample)
+ end = sample[-1]
+ output.append(
+ f"C {control_a[0]:.2f} {control_a[1]:.2f} "
+ f"{control_b[0]:.2f} {control_b[1]:.2f} {end[0]:.2f} {end[1]:.2f}"
+ )
+ return " ".join(output)
+
+
+def _mask_distance(mask: np.ndarray) -> np.ndarray:
+ """Two-pass chamfer distance to the background in mask-pixel units."""
+ distance = np.where(mask, np.inf, 0.0).astype(np.float64)
+ diagonal = math.sqrt(2.0)
+ for y in range(distance.shape[0]):
+ for x in range(distance.shape[1]):
+ if not mask[y, x]:
+ continue
+ candidates = []
+ if y:
+ candidates.append(distance[y - 1, x] + 1)
+ if x:
+ candidates.append(distance[y - 1, x - 1] + diagonal)
+ if x + 1 < distance.shape[1]:
+ candidates.append(distance[y - 1, x + 1] + diagonal)
+ if x:
+ candidates.append(distance[y, x - 1] + 1)
+ distance[y, x] = min(candidates, default=distance[y, x])
+ for y in range(distance.shape[0] - 1, -1, -1):
+ for x in range(distance.shape[1] - 1, -1, -1):
+ if not mask[y, x]:
+ continue
+ candidates = [distance[y, x]]
+ if y + 1 < distance.shape[0]:
+ candidates.append(distance[y + 1, x] + 1)
+ if x:
+ candidates.append(distance[y + 1, x - 1] + diagonal)
+ if x + 1 < distance.shape[1]:
+ candidates.append(distance[y + 1, x + 1] + diagonal)
+ if x + 1 < distance.shape[1]:
+ candidates.append(distance[y, x + 1] + 1)
+ distance[y, x] = min(candidates)
+ return distance
+
+
+def _trace_sections(trace: np.ndarray, segments: int) -> list[np.ndarray]:
+ count = max(1, min(segments, math.ceil((len(trace) - 1) / 8)))
+ boundaries = np.linspace(0, len(trace) - 1, count + 1, dtype=int)
+ return [trace[first : last + 1] for first, last in itertools.pairwise(boundaries)]
+
+
+def mask_stroke(
+ mask: np.ndarray, *, segments: int = 8, overlap_pixels: int = 0
+) -> tuple[str, float] | None:
+ """Return a conservative centreline stroke for one thin mask component.
+
+ SAMVG itself uses closed filled shapes. This optional hybrid extension is
+ deliberately strict: a component must be long, narrow, and have no holes
+ before it can be represented by a stroke. Other masks preserve SAMVG's
+ original filled-path treatment.
+ """
+ if overlap_pixels:
+ mask = _binary_dilation(mask, overlap_pixels)
+ _ys, xs = np.nonzero(mask)
+ if len(xs) < 8:
+ return None
+ # A hole is topology that a single centreline cannot preserve.
+ if len(_loops(mask)) != 1:
+ return None
+
+ traces = _skeleton_traces(mask)
+ if len(traces) != 1:
+ return None
+ trace = traces[0]
+ length = float(np.linalg.norm(np.diff(trace, axis=0), axis=1).sum())
+ # Arc length, rather than a bounding-box axis, preserves strongly curved
+ # thin components whose width and height are similar.
+ estimated_width = len(xs) / max(length, 1.0)
+ if length < 12 or estimated_width > min(8.0, length * 0.3):
+ return None
+ distance = _mask_distance(mask)
+ widths = [2 * (distance[int(y), int(x)] - 0.5) for x, y in trace]
+ data = _trace_path_data(trace, segments)
+ return data, max(1.0, float(np.median(widths)))
+
+
+def mask_strokes(
+ mask: np.ndarray, *, segments: int = 8, overlap_pixels: int = 0
+) -> list[tuple[str, float]]:
+ """Trace a thin component into independently editable constant-width paths.
+
+ A branch becomes one path per medial-axis edge. Each long edge is divided
+ into cubic sections and each section gets its local median width, giving an
+ SVG approximation of a variable-width centreline without nonstandard SVG
+ extensions. Round caps and joins make the adjacent sections continuous.
+ """
+ if overlap_pixels:
+ mask = _binary_dilation(mask, overlap_pixels)
+ _ys, xs = np.nonzero(mask)
+ if len(xs) < 8:
+ return []
+ if len(_loops(mask)) != 1:
+ return []
+ traces = _skeleton_traces(mask)
+ length = sum(
+ float(np.linalg.norm(np.diff(trace, axis=0), axis=1).sum()) for trace in traces
+ )
+ estimated_width = len(xs) / max(length, 1.0)
+ if length < 12 or estimated_width > min(8.0, length * 0.3):
+ return []
+ distance = _mask_distance(mask)
+ output: list[tuple[str, float]] = []
+ for trace in traces:
+ for section in _trace_sections(trace, segments):
+ if len(section) < 2:
+ continue
+ widths = [2 * (distance[int(y), int(x)] - 0.5) for x, y in section]
+ output.append(
+ (_trace_path_data(section, 1), max(1.0, float(np.median(widths))))
+ )
+ return output
+
+
+def _layer_svg_attributes(
+ layer: MaskLayer, segments: int, *, hybrid_strokes: bool = True
+) -> 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}"
+ strokes = (
+ mask_strokes(layer.mask, segments=segments, overlap_pixels=layer.overlap_pixels)
+ if hybrid_strokes
+ else []
+ )
+ if strokes:
+ return [
+ {
+ "d": data,
+ "fill": "none",
+ "stroke": colour,
+ "stroke-width": f"{width:.2f}",
+ "stroke-linecap": "round",
+ "stroke-linejoin": "round",
+ }
+ for data, width in strokes
+ ]
+ data = mask_path(layer.mask, segments=segments, overlap_pixels=layer.overlap_pixels)
+ if data is None:
+ return []
+ return [{"d": data, "fill": colour, "fill-rule": "evenodd"}]
+
+
def generate_svg(
image: Image.Image,
masks: list[np.ndarray] | None = None,
@@ -557,6 +1388,10 @@ def generate_svg(
max_layers: int = 512,
segments: int = 16,
fill_holes: bool = True,
+ hybrid_strokes: bool = True,
+ ocr: bool = True,
+ max_side: int | None = SAMVG_MAX_SIDE,
+ rasterize: Callable[[str, int, int], bytes] | None = None,
) -> str:
"""Generate SAMVG's traced, pre-optimisation SVG from a target image."""
image = image.convert("RGB")
@@ -576,21 +1411,25 @@ def generate_svg(
min_impact=min_impact,
max_layers=max_layers,
fill_holes=fill_holes,
+ max_side=max_side,
)
)
paths = []
for layer in layers:
- data = mask_path(
- layer.mask, segments=segments, overlap_pixels=layer.overlap_pixels
- )
- if data:
- colour = f"#{layer.colour[0]:02x}{layer.colour[1]:02x}{layer.colour[2]:02x}"
- paths.append(f'')
+ for attributes in _layer_svg_attributes(
+ layer, segments, hybrid_strokes=hybrid_strokes
+ ):
+ markup = " ".join(f'{key}="{value}"' for key, value in attributes.items())
+ paths.append(f"")
width, height = image.size
- return (
+ svg = (
f'"
)
+ text_layers = detect_text(image) if ocr and masks is None else []
+ if text_layers and rasterize is not None:
+ return _accept_text_layers(svg, image, text_layers, rasterize)
+ return _append_text_layers(svg, text_layers)
def residual_prompt_points(
@@ -602,8 +1441,8 @@ def residual_prompt_points(
max_points: int = 16,
) -> list[tuple[int, int]]:
"""Locate SAMVG's convolved, thresholded residual components."""
- from scipy.ndimage import label
- from scipy.signal import fftconvolve
+ 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
@@ -617,8 +1456,11 @@ def residual_prompt_points(
# 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 = fftconvolve(padded, kernel / kernel.sum(), mode="valid")
- labels, count = label(smoothed >= threshold)
+ 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)
@@ -629,25 +1471,37 @@ def residual_prompt_points(
return [(x, y) for _score, x, y in sorted(points, reverse=True)[:max_points]]
-def _append_layers(svg: str, layers: list[MaskLayer], segments: int) -> str:
+def _append_layers(
+ svg: str,
+ layers: list[MaskLayer],
+ segments: int,
+ *,
+ hybrid_strokes: bool = True,
+) -> str:
"""Add newly prompted paths to an already optimised SVG."""
root = ET.fromstring(svg)
for layer in layers:
- data = mask_path(
- layer.mask, segments=segments, overlap_pixels=layer.overlap_pixels
- )
- if not data:
- continue
- colour = f"#{layer.colour[0]:02x}{layer.colour[1]:02x}{layer.colour[2]:02x}"
- ET.SubElement(
- root,
- "{http://www.w3.org/2000/svg}path",
- {
- "d": data,
- "fill": colour,
- "fill-rule": "evenodd",
- },
+ for attributes in _layer_svg_attributes(
+ layer, segments, hybrid_strokes=hybrid_strokes
+ ):
+ ET.SubElement(
+ root,
+ "{http://www.w3.org/2000/svg}path",
+ attributes,
+ )
+ return ET.tostring(root, encoding="unicode")
+
+
+def _append_text_layers(svg: str, layers: list[TextLayer]) -> str:
+ """Append editable OCR text without changing the pre-existing drawing."""
+ if not layers:
+ return svg
+ root = ET.fromstring(svg)
+ for layer in layers:
+ element = ET.SubElement(
+ root, "{http://www.w3.org/2000/svg}text", _text_svg_attributes(layer)
)
+ element.text = layer.text
return ET.tostring(root, encoding="unicode")
@@ -663,14 +1517,52 @@ def _mse(image: Image.Image, rendered: Image.Image) -> float:
return float(((target - candidate) ** 2).mean())
+def _text_error_tolerance(layer: TextLayer, image: Image.Image) -> float:
+ """Return the whole-image MSE budget for this one text bounding box."""
+ padding = 2
+ width = min(image.width, max(1, math.ceil(layer.width) + padding * 2))
+ height = min(image.height, max(1, math.ceil(layer.height) + padding * 2))
+ affected_fraction = (width * height) / (image.width * image.height)
+ return affected_fraction * (255 * OCR_TEXT_RMSE_TOLERANCE) ** 2
+
+
+def _accept_text_layers(
+ svg: str,
+ image: Image.Image,
+ layers: list[TextLayer],
+ rasterize: Callable[[str, int, int], bytes],
+) -> str:
+ """Retain OCR text that improves, or only negligibly worsens, pixel loss.
+
+ A VLM's asserted confidence is not evidence that a word is present. The
+ same rasterisation used to score the seed is the final verifier, including
+ font mismatch, positioning, and any existing SAM paths beneath the text.
+ """
+ accepted = svg
+ error = _mse(image, _render_svg(accepted, image, rasterize))
+ retained = 0
+ for layer in layers:
+ candidate = _append_text_layers(accepted, [layer])
+ candidate_error = _mse(image, _render_svg(candidate, image, rasterize))
+ if candidate_error <= error + _text_error_tolerance(layer, image):
+ accepted, error = candidate, candidate_error
+ retained += 1
+ log.info(
+ "SAMVG OCR: retained %d/%d text layer(s) after pixel verification.",
+ retained,
+ len(layers),
+ )
+ return accepted
+
+
def _accepted_fit(
svg: str, image: Image.Image, *, rasterize, steps: int
) -> tuple[str, Image.Image]:
"""Keep a differentiable fit only when the actual SVG renderer improves."""
- from vectrify.refine.paths import fit_filled_svg
+ from vectrify.refine.paths import fit_filled_svg_bounded
before = _render_svg(svg, image, rasterize)
- fitted = fit_filled_svg(svg, image, steps=steps)
+ fitted = fit_filled_svg_bounded(svg, image, rasterize=rasterize, steps=steps)
after = _render_svg(fitted, image, rasterize)
if _mse(image, after) <= _mse(image, before):
return fitted, after
@@ -686,7 +1578,8 @@ def vectorize_svg(
min_pixels: int = 32,
min_impact: float = 1e-5,
max_layers: int = 512,
- segments: int = 8,
+ segments: int = 16,
+ max_side: int | None = SAMVG_MAX_SIDE,
) -> str:
"""Run SAMVG's two 500-step optimise-and-recover phases.
@@ -695,38 +1588,56 @@ def vectorize_svg(
built-in filled-path optimiser so SAMVG has no external renderer dependency.
"""
image = image.convert("RGB")
- layers = retrieve_layers(
- image,
- min_pixels=min_pixels,
- min_impact=min_impact,
- max_layers=max_layers,
- )
- initial = _append_layers(
- f'',
- layers,
- segments,
- )
- first, first_render = _accepted_fit(
- initial, image, rasterize=rasterize, steps=steps
- )
- 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),
- existing=layers,
- initial_canvas=np.asarray(first_render, dtype=np.uint8),
- initial_coverage=coverage,
- min_pixels=min_pixels,
- min_impact=min_impact,
- max_layers=max_layers,
- )[len(layers) :]
- log.info(
- "SAMVG residual pass: %d prompt(s), %d accepted added path(s).",
- len(points),
- len(added),
- )
- return _accepted_fit(
- _append_layers(first, added, segments), image, rasterize=rasterize, steps=steps
- )[0]
+ runtime = _sam_runtime()
+ try:
+ layers = retrieve_layers(
+ image,
+ min_pixels=min_pixels,
+ min_impact=min_impact,
+ max_layers=max_layers,
+ max_side=max_side,
+ _runtime=runtime,
+ )
+ initial = _append_layers(
+ f'',
+ layers,
+ segments,
+ hybrid_strokes=False,
+ )
+ first, first_render = _accepted_fit(
+ initial, image, rasterize=rasterize, steps=steps
+ )
+ 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,
+ min_pixels=min_pixels,
+ min_impact=min_impact,
+ max_layers=max_layers,
+ )[len(layers) :]
+ log.info(
+ "SAMVG residual pass: %d prompt(s), %d accepted added path(s).",
+ len(points),
+ len(added),
+ )
+ return _accepted_fit(
+ _append_layers(first, added, segments, hybrid_strokes=False),
+ image,
+ rasterize=rasterize,
+ steps=steps,
+ )[0]
+ finally:
+ del runtime
+ try:
+ import torch
+
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ except ImportError: # pragma: no cover - installation-specific
+ pass
diff --git a/src/vectrify/vector/runner.py b/src/vectrify/vector/runner.py
index fbee3425..a6559758 100644
--- a/src/vectrify/vector/runner.py
+++ b/src/vectrify/vector/runner.py
@@ -478,7 +478,14 @@ def run_vector_search(
log.info("SAMVG-inspired seed skipped: it is available for SVG only.")
else:
try:
- content = format_plugin.extract_from_llm(generate_svg(original_img))
+ content = format_plugin.extract_from_llm(
+ generate_svg(
+ original_img,
+ rasterize=lambda svg, width, height: format_plugin.rasterize(
+ svg, out_w=width, out_h=height
+ ),
+ )
+ )
valid, error = format_plugin.validate(content)
if not valid:
raise ValueError(error or "generated SVG failed validation")
diff --git a/tests/formats/svg/test_plugin.py b/tests/formats/svg/test_plugin.py
index 65edf02f..68928138 100644
--- a/tests/formats/svg/test_plugin.py
+++ b/tests/formats/svg/test_plugin.py
@@ -222,6 +222,11 @@ def test_a_usable_reply_is_still_applied():
''
)
+_FILLED = (
+ ''
+)
+
def test_the_path_fit_is_only_offered_where_it_is_cheap_enough():
"""It costs ~0.5s on a GPU against ~9s on one CPU thread, where an ordinary
@@ -258,6 +263,34 @@ def test_a_drawing_with_nothing_fittable_reports_a_blank_draw():
assert origin == PATH_FIT
+def test_path_fit_dispatches_opaque_fills_to_the_samvg_renderer(monkeypatch):
+ """SAMVG-style fill seeds must not fall back to the sampled stroke fit."""
+ from vectrify.formats.svg import plugin as plugin_module
+ from vectrify.refine.paths import PATH_FIT
+
+ plugin = SvgPlugin()
+ seen = {}
+
+ def fit(svg, reference_png, *, rasterize, weights, gpu_gate):
+ seen["svg"] = svg
+ seen["reference"] = reference_png
+ seen["gpu_gate"] = gpu_gate
+ assert rasterize(svg, 64, 64)
+ assert weights is None
+ return svg.replace("#111111", "#ff0000")
+
+ monkeypatch.setattr(plugin_module, "fit_available", lambda: True)
+ monkeypatch.setattr(plugin_module, "fit_svg_primitives_locally", fit)
+ reference = plugin.rasterize(_FILLED, 64, 64)
+
+ content, origin = plugin.mutate(_FILLED, operator=PATH_FIT, reference_png=reference)
+
+ assert origin == PATH_FIT
+ assert "#ff0000" in content
+ assert seen["svg"] == _FILLED
+ assert seen["reference"] == reference
+
+
def test_a_width_on_the_path_is_found_as_readily_as_one_on_the_group():
"""The width can be declared on the element, a group above it, or the root.
One model wrote it on every path and none on their groups; a lookup that
@@ -317,50 +350,49 @@ def test_a_full_device_skips_the_fit_instead_of_failing_the_task():
Every worker that fits holds a context of a few hundred MB and there are as
many workers as cores, so running out is a normal condition, not a bug.
"""
- import vectrify.refine.paths as paths
+ import vectrify.formats.svg.plugin as plugin_module
from vectrify.refine.paths import PATH_FIT
plugin = SvgPlugin()
png = plugin.rasterize(_STROKED, 64, 64)
- original = paths.fit_random_group
def out_of_memory(*_args, **_kwargs):
raise RuntimeError("CUDA error: out of memory")
- from vectrify.formats.svg import plugin as plugin_module
-
- plugin_module.fit_random_group = out_of_memory
+ original = plugin_module.fit_svg_primitives_locally
+ original_available = plugin_module.fit_available
+ plugin_module.fit_svg_primitives_locally = out_of_memory
plugin_module.fit_available = lambda: True
try:
content, origin = plugin.mutate(_STROKED, operator=PATH_FIT, reference_png=png)
finally:
- plugin_module.fit_random_group = original
- plugin_module.fit_available = paths.fit_available
+ plugin_module.fit_svg_primitives_locally = original
+ plugin_module.fit_available = original_available
assert content == _STROKED
assert origin == PATH_FIT
def test_an_unrelated_failure_in_the_fit_is_not_swallowed():
- import vectrify.refine.paths as paths
from vectrify.formats.svg import plugin as plugin_module
from vectrify.refine.paths import PATH_FIT
plugin = SvgPlugin()
png = plugin.rasterize(_STROKED, 64, 64)
- original = plugin_module.fit_random_group
+ original = plugin_module.fit_svg_primitives_locally
+ original_available = plugin_module.fit_available
def bug(*_args, **_kwargs):
raise ValueError("something genuinely wrong")
- plugin_module.fit_random_group = bug
+ plugin_module.fit_svg_primitives_locally = bug
plugin_module.fit_available = lambda: True
try:
with pytest.raises(ValueError, match="genuinely wrong"):
plugin.mutate(_STROKED, operator=PATH_FIT, reference_png=png)
finally:
- plugin_module.fit_random_group = original
- plugin_module.fit_available = paths.fit_available
+ plugin_module.fit_svg_primitives_locally = original
+ plugin_module.fit_available = original_available
def test_path_fit_receives_the_shared_gpu_gate():
@@ -373,7 +405,7 @@ class Gate:
plugin = SvgPlugin()
plugin.gpu_gate = gate = Gate()
png = plugin.rasterize(_STROKED, 64, 64)
- original_fit = plugin_module.fit_random_group
+ original_fit = plugin_module.fit_svg_primitives_locally
original_available = plugin_module.fit_available
seen = {}
@@ -381,12 +413,12 @@ def fit(*args, **kwargs):
seen["gpu_gate"] = kwargs["gpu_gate"]
return args[0]
- plugin_module.fit_random_group = fit
+ plugin_module.fit_svg_primitives_locally = fit
plugin_module.fit_available = lambda: True
try:
plugin.mutate(_STROKED, operator=PATH_FIT, reference_png=png)
finally:
- plugin_module.fit_random_group = original_fit
+ plugin_module.fit_svg_primitives_locally = original_fit
plugin_module.fit_available = original_available
assert seen["gpu_gate"] is gate
diff --git a/tests/refine/test_fidelity.py b/tests/refine/test_fidelity.py
index e17e96c5..9c45648d 100644
--- a/tests/refine/test_fidelity.py
+++ b/tests/refine/test_fidelity.py
@@ -138,7 +138,7 @@ def test_a_pinned_vertex_does_not_move():
).convert("L")
path_d = "M 300 300 C 340 280 400 280 440 300"
knots = to_knots(parse_cubics(path_d))
- fitted, _first, _last = fit_group(
+ fitted, _widths, _colours, _first, _last = fit_group(
[path_d], [3.5], target, target, steps=6, pinned={0}
)
moved = to_knots(parse_cubics(fitted[0]))
diff --git a/tests/refine/test_filled_paths.py b/tests/refine/test_filled_paths.py
index a9c5889d..9428fcda 100644
--- a/tests/refine/test_filled_paths.py
+++ b/tests/refine/test_filled_paths.py
@@ -1,4 +1,5 @@
import io
+import xml.etree.ElementTree as ET
import numpy as np
import pytest
@@ -17,6 +18,7 @@
_tiled_large_path_coverage,
_xing_loss,
fit_filled_svg,
+ fit_opaque_fills_locally,
parse_filled_cubics,
)
@@ -127,9 +129,11 @@ def test_native_even_odd_coverage_stays_cairo_validated():
torch.tensor(contour, dtype=torch.float32, device="cuda")
for contour in parse_filled_cubics(DONUT_PATH)
]
- native = _fill_path_coverage(
- contours, (0, 0, size, size), fill_rule="evenodd"
- ).cpu().numpy()
+ native = (
+ _fill_path_coverage(contours, (0, 0, size, size), fill_rule="evenodd")
+ .cpu()
+ .numpy()
+ )
head = f'"
drawn = f'{head}'
@@ -167,9 +171,7 @@ def test_native_analytic_cubic_coverage_stays_cairo_validated():
dtype=torch.float32,
device="cuda",
)
- contour = torch.cat((contour, contour[:1].expand(16 - len(contour), -1, -1)))[
- None
- ]
+ contour = torch.cat((contour, contour[:1].expand(16 - len(contour), -1, -1)))[None]
native = _fill_coverages(contour, (0, 0, size, size), subpixels=4).cpu().numpy()
head = f'"
@@ -207,9 +209,7 @@ def test_native_analytic_multi_contour_coverage_preserves_a_hole():
]
controls = torch.cat(
[
- torch.cat((contour, contour[:1].expand(16 - len(contour), -1, -1)))[
- None
- ]
+ torch.cat((contour, contour[:1].expand(16 - len(contour), -1, -1)))[None]
for contour in contours
]
).requires_grad_()
@@ -222,6 +222,7 @@ def test_native_analytic_multi_contour_coverage_preserves_a_hole():
assert controls.grad is not None
assert controls.grad.abs().sum() > 0
+
SVG = (
'