diff --git a/AGENTS.md b/AGENTS.md index 78c03eaf..2a9008fe 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -42,7 +42,7 @@ Releases: follow [the release procedure](docs/releasing.md) and `scripts/Publish - `OpenNest.IO`: ACadSharp import/export and ZIP-based `.nest` persistence; depends on Core only (an IO.Tests guard enforces it). Best fits persist only when callers pass `BestFitCacheStore.Shared` as `NestReader`/`NestWriter.BestFitStore`. All DXF-to-Drawing conversion goes through `CadImporter`: `Import` + `BuildDrawing` for editable/reporting flows, `ImportDrawing` for headless callers. Preserve source offsets, entity IDs, suppressed entities and bends. Bend repair is opt-in, requires explicit source units and may not alter cut geometry or unrelated marks. - `OpenNest`: WinForms UI (`Forms/`, `Controls/PlateView`, `Actions/`). `OpenNest.Data` holds cross-platform persistence; new-nest defaults live in `%APPDATA%\OpenNest\defaults.json` (in Database mode, the server's shared copy once it holds one). WinForms `Properties.Settings` persist to `%APPDATA%\OpenNest\user-settings.json` through `AppDataSettingsProvider`, never inside the Velopack install folder. Posts live in `Posts/OpenNest.Posts./` and deploy to the desktop output's `Posts/` directory. Desktop nest edits go through `BeginChange`/`RecordChange`; see [undo and redo](docs/undo-redo.md). - `OpenNest.Desktop` (net8.0): toolkit-neutral desktop logic shared by the WinForms app and `OpenNest.Avalonia`, the cross-platform desktop in progress. It must not reference WinForms/WindowsDesktop, `System.Drawing.Common` or Avalonia (`OpenNest.Desktop.Tests` guards this, including transitive packages); put toolkit-specific rendering, dialogs and settings behind adapters in the front ends. -- `OpenNest.Console`, `OpenNest.Mcp`, `OpenNest.Api`: front ends; `OpenNest.Benchmark`: whole-job engine comparisons; `OpenNest.Gpu`: GPU evaluators; `OpenNest.Training`: ML data collection. Benchmark timing comparisons require `--parallel 1`; validate layouts and fulfillment, not just elapsed time. +- `OpenNest.Console`, `OpenNest.Mcp`, `OpenNest.Api`: front ends; `OpenNest.Benchmark`: whole-job engine comparisons; `OpenNest.Gpu`: GPU evaluators; `OpenNest.Training`: ML data collection and the Python sketch-search prototype (`sketch-search/README.md`). Benchmark timing comparisons require `--parallel 1`; validate layouts and fulfillment, not just elapsed time. ## Geometry and ownership safeguards diff --git a/OpenNest.Training/sketch-search/.gitignore b/OpenNest.Training/sketch-search/.gitignore new file mode 100644 index 00000000..c18dd8d8 --- /dev/null +++ b/OpenNest.Training/sketch-search/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/OpenNest.Training/sketch-search/README.md b/OpenNest.Training/sketch-search/README.md new file mode 100644 index 00000000..7a0a55e9 --- /dev/null +++ b/OpenNest.Training/sketch-search/README.md @@ -0,0 +1,102 @@ +# Sketch search (prototype) + +Finds drawings that look like a hand sketch, a phone photo of one, or an MS Paint +outline. A DINOv2 image model compares the sketch with a rendering of every +drawing. The model is pretrained, optionally fine-tuned on synthetic sketches of +the library itself (see [Fine-tune](#fine-tune)). + +## Setup + +```sh +pip install torch --index-url https://download.pytorch.org/whl/cu124 +pip install -r requirements.txt +``` + +A CPU-only torch works, but embedding is much slower than on a GPU. + +## Build the index + +```sh +python build_index.py --index C:\SketchIndex --server http://barge.lan:8091 +python build_index.py --index C:\SketchIndex --dxf X:\Archive +``` + +`--server` reads the server's drawing library with GET requests only. `--dxf` +walks a folder of DXF files and skips text, dimensions, hatches and +`ETCH`/`SCRIBE` layers. Both options can be repeated and combined. Programs are +cached under `/cache`. Re-running the build only processes drawings whose +revision or file changed. + +## Search + +```sh +python search.py --index C:\SketchIndex sketch.png --open +python search.py --index C:\SketchIndex photo.jpg --photo --size 10x4.5 +``` + +- **MS Paint:** draw the outline in a dark color and save the file, or crop a + screenshot to the canvas. Window chrome in the image counts as part of the + sketch. Either outlines or a filled silhouette work. +- **Photos and scans:** pass `--photo` to use local thresholding, which handles + uneven lighting. Crop to the sketch and take the photo square-on. +- **`--size WxL`:** keeps only drawings of that overall size (either + orientation, inches or millimetres, default ±15%). The size is a hard filter, + so use it only when the sketch's dimensions are reliable. +- **Leave out text and small details:** draw the outline and the holes. + Rotation and mirroring do not matter. + +Results go to `/last-search.html`. The page shows the sketch as the model +sees it, followed by the closest drawings with name, customer, material and +size. + +## Measure accuracy + +```sh +python evaluate.py --index C:\SketchIndex --synthetic 500 +python evaluate.py --index C:\SketchIndex --real C:\Sketches +``` + +`--synthetic` turns random held-out drawings into fake sketches and reports how +often the source drawing ranks in the top N. Fake sketches have wobbly lines, +±15% proportions, dropped small holes, random rotation, mirroring, line weight +and fill. "Same shape" also counts a drawing that renders identically to the +answer, such as duplicates and unchanged revisions. + +`--real` takes a folder of real sketches, each named after its drawing +(`BRKT-12.png`, or `BRKT-12__2.jpg` for extra sketches of the same part). Real +sketches are the result that decides whether the approach works. + +## How it works + +1. Each drawing's cut geometry is rasterized: G-code from the server, or DXF + entities. Scribe, etch and break motions are skipped. +2. Drawings and sketches share one `normalize()` step. It reduces the ink to + 1px centre lines (a filled silhouette becomes its boundary), crops and scales + the result into a 224×224 canvas, and redraws it at a fixed stroke width. +3. DINOv2 embeds the image: the CLS token joined with the mean patch token. +4. The query is embedded in all 8 rotations and mirrors, and each drawing scores + its best match. A small penalty applies when the drawing's proportions + differ from the sketch's (`ASPECT_WEIGHT` in `search.py`). + +## Fine-tune + +The pretrained model alone is weak on a 40k-drawing library. Training it on +synthetic sketches of the library itself needs no labelling: + +```sh +python train.py --index C:\SketchIndex --out C:\SketchIndex\model +python build_index.py --index C:\SketchIndex --server http://barge.lan:8091 --model C:\SketchIndex\model +python evaluate.py --index C:\SketchIndex --synthetic 400 +``` + +Each epoch pairs every training drawing's rendering with a freshly generated +fake sketch (`synthetic.py`, kept in the drawing's orientation because search +tries all 8). The model learns to pick each sketch's own drawing out of the +batch. Two groups of drawings are left out of training: + +- Near-duplicate renderings, which would count as wrong answers within a batch. +- About 10% of drawings chosen by key (`is_holdout`). `evaluate.py` samples only + these, so its scores are not inflated by training. + +Rebuilding the index with a new `--model` reuses the existing renders and only +re-embeds them. diff --git a/OpenNest.Training/sketch-search/build_index.py b/OpenNest.Training/sketch-search/build_index.py new file mode 100644 index 00000000..19d2a143 --- /dev/null +++ b/OpenNest.Training/sketch-search/build_index.py @@ -0,0 +1,191 @@ +"""Builds or refreshes a sketch-search index from an OpenNest server and/or DXF folders. + + python build_index.py --index C:\\SketchIndex --server http://barge.lan:8091 + python build_index.py --index C:\\SketchIndex --dxf X:\\Archive + +Re-running only downloads, renders and embeds drawings whose revision changed. +""" +import argparse +import concurrent.futures as cf +import io +import json +import os +import sys +import time +import urllib.request +import zipfile +from pathlib import Path + +import numpy as np +from PIL import Image + +import sketch_common as sc + + +def list_server(server, limit=None): + items, offset = [], 0 + while True: + page = json.load(urllib.request.urlopen(f"{server}/api/drawings/query?limit=500&offset={offset}", timeout=60)) + for r in page["items"]: + items.append({ + "key": f"server:{r['id']}", "source": "server", "id": r["id"], "revision": r["revision"], + "sha": r["packageSha256"], "name": r["name"], "customer": r.get("customer", ""), + "material": " ".join(x for x in (r.get("materialName"), r.get("materialGrade")) if x), + "thickness": r.get("thickness"), "units": r.get("units", ""), + }) + offset += len(page["items"]) + if not page["items"] or offset >= page["total"] or (limit and offset >= limit): + return items[:limit] if limit else items + + +def list_dxf(folder, limit=None): + items = [] + for root, _, files in os.walk(folder): + for f in files: + if f.lower().endswith(".dxf"): + p = Path(root, f) + st = p.stat() + items.append({"key": f"dxf:{p}", "source": "dxf", "path": str(p), "name": p.stem, + "sha": f"{st.st_size}-{int(st.st_mtime)}", "units": ""}) + if limit and len(items) >= limit: + return items + return items + + +def program_path(index, item): + return Path(index, "cache", f"{item['id']}-r{item['revision']}.gcode") + + +def download(server, index, item): + path = program_path(index, item) + if path.exists(): + return None + for attempt in range(3): + try: + data = urllib.request.urlopen(f"{server}/api/drawings/{item['id']}/file", timeout=60).read() + z = zipfile.ZipFile(io.BytesIO(data)) + programs = sorted(n for n in z.namelist() if n.startswith("programs/program-") and not n.endswith("-subs")) + path.write_bytes(z.read(programs[0]) if programs else b"") + return None + except Exception as e: # noqa: BLE001 - reported per drawing + err = e + time.sleep(1 + attempt) + return f"{item['name']}: {err}" + + +def load_polys(index, item): + if item["source"] == "server": + return sc.parse_program(program_path(index, item).read_text(errors="replace")) + return sc.read_dxf(item["path"]) + + +def render_one(args): + index, item = args + try: + polys = load_polys(index, item) + if not polys: + return item["key"], None, "no cut geometry" + lo, hi = sc.bounds(polys) + sc.normalize(sc.render_mask(polys)).convert("L").save(Path(index, "thumbs", thumb_name(item))) + return item["key"], (float(hi[0] - lo[0]), float(hi[1] - lo[1])), None + except Exception as e: # noqa: BLE001 - reported per drawing + return item["key"], None, str(e) + + +def thumb_name(item): + if item["source"] == "server": + return f"{item['id']}.png" + import hashlib + return "dxf-" + hashlib.sha1(item["path"].encode()).hexdigest()[:16] + ".png" + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--index", required=True, help="index folder (created if missing)") + ap.add_argument("--server", action="append", default=[], help="OpenNest server URL") + ap.add_argument("--dxf", action="append", default=[], help="folder of DXF files (searched recursively)") + ap.add_argument("--model", default="facebook/dinov2-small", help="Hugging Face model or fine-tuned model folder") + ap.add_argument("--limit", type=int, help="index at most this many drawings per source (testing)") + ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 2)) + args = ap.parse_args() + if not args.server and not args.dxf: + ap.error("give at least one --server or --dxf source") + + index = Path(args.index) + for sub in ("cache", "thumbs"): + Path(index, sub).mkdir(parents=True, exist_ok=True) + + items = [] + for server in args.server: + server = server.rstrip("/") + found = list_server(server, args.limit) + print(f"{server}: {len(found)} drawings") + with cf.ThreadPoolExecutor(12) as pool: + errors = [e for e in pool.map(lambda it: download(server, index, it), found) if e] + for e in errors[:10]: + print(" download failed:", e) + items += found + for folder in args.dxf: + found = list_dxf(folder, args.limit) + print(f"{folder}: {len(found)} DXF files") + items += found + + config = {"model": args.model} + old_items, old_emb = {}, None + if Path(index, "items.jsonl").exists() and Path(index, "config.json").exists(): + old_list = [json.loads(l) for l in Path(index, "items.jsonl").read_text(encoding="utf-8").splitlines()] + old_items = {it["key"]: (it, i) for i, it in enumerate(old_list)} + if json.loads(Path(index, "config.json").read_text()) == config: + old_emb = np.load(Path(index, "embeddings.npy")) + + # Renders survive a model change; embeddings only survive with the same model. + embedded, unembedded, todo = [], [], [] + for it in items: + prev = old_items.get(it["key"]) + if prev and prev[0]["sha"] == it["sha"] and Path(index, "thumbs", thumb_name(it)).exists(): + it.update(width=prev[0]["width"], height=prev[0]["height"]) + if old_emb is not None: + embedded.append((it, old_emb[prev[1]])) + else: + unembedded.append(it) + else: + todo.append(it) + print(f"reusing {len(embedded) + len(unembedded)} renders, rendering {len(todo)}") + + by_key = {it["key"]: it for it in todo} + failures = [] + started = time.time() + with cf.ProcessPoolExecutor(args.workers) as pool: + for n, (key, size, err) in enumerate(pool.map(render_one, [(str(index), it) for it in todo], chunksize=16), 1): + if err: + failures.append((key, err)) + else: + by_key[key].update(width=size[0], height=size[1]) + unembedded.append(by_key[key]) + if n % 2000 == 0: + print(f" rendered {n}/{len(todo)} ({time.time() - started:.0f}s)") + print(f"rendered {len(todo) - len(failures)} in {time.time() - started:.0f}s, {len(failures)} skipped") + for key, err in failures[:15]: + print(" skipped", key, "-", err) + + rows = [e for _, e in embedded] + if unembedded: + embedder = sc.Embedder(args.model) + started = time.time() + load = lambda it: Image.open(Path(index, "thumbs", thumb_name(it))).convert("RGB") + with cf.ThreadPoolExecutor(8) as pool: + for i in range(0, len(unembedded), 2048): + rows += list(embedder.embed(list(pool.map(load, unembedded[i:i + 2048])))) + print(f"embedded {len(unembedded)} in {time.time() - started:.0f}s") + + all_items = [it for it, _ in embedded] + unembedded + for it in all_items: + it["thumb"] = thumb_name(it) + np.save(Path(index, "embeddings.npy"), np.stack(rows).astype(np.float16)) + Path(index, "items.jsonl").write_text("\n".join(json.dumps(it) for it in all_items), encoding="utf-8") + Path(index, "config.json").write_text(json.dumps(config)) + print(f"index: {len(all_items)} drawings in {index}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/OpenNest.Training/sketch-search/evaluate.py b/OpenNest.Training/sketch-search/evaluate.py new file mode 100644 index 00000000..e9ce3ac8 --- /dev/null +++ b/OpenNest.Training/sketch-search/evaluate.py @@ -0,0 +1,117 @@ +"""Measures how often search finds the right drawing. + +Synthetic: turns random held-out drawings (never trained on) into fake hand +sketches and records where the source drawing ranks among the whole index. + + python evaluate.py --index C:\\SketchIndex --synthetic 300 + +Real: a folder of sketch images, each named after the drawing it should find +(the drawing name, or name + anything after "__", e.g. "BRKT-12__photo.jpg"). + + python evaluate.py --index C:\\SketchIndex --real C:\\Sketches [--photo] + +A "same shape" hit also counts a drawing whose rendering is effectively identical +to the answer (duplicates and unchanged revisions are common in a part library). +""" +import argparse +import json +import random +import sys +from pathlib import Path + +import numpy as np +from PIL import Image + +import build_index +from synthetic import sketchify +import sketch_common as sc +from search import ASPECT_WEIGHT, SketchIndex + +SAME_SHAPE = 0.985 +CUTOFFS = (1, 5, 10, 20, 50) + + +def rank_of(index, results_keys, truth_idx): + """(strict rank, same-shape rank), 1-based; None when beyond the list.""" + truth_emb = index.embeddings[truth_idx] + strict = same = None + for r, i in enumerate(results_keys, 1): + if strict is None and i == truth_idx: + strict = r + if same is None and (i == truth_idx or float(index.embeddings[i] @ truth_emb) >= SAME_SHAPE): + same = r + return strict, same + + +def full_ranking(index, mask, k, aspect_weight): + scores = index.scores(sc.normalize(mask), sc.aspect(mask), aspect_weight) + return list(np.argsort(-scores)[:k]) + + +def report(name, ranks, k): + print(f"\n{name}: {len(ranks)} queries") + for label, col in (("exact drawing", 0), ("same shape", 1)): + found = [r[col] for r in ranks] + cells = " ".join(f"top{c}: {100 * sum(1 for f in found if f and f <= c) / len(found):5.1f}%" for c in CUTOFFS) + print(f" {label:13} {cells} (median rank {np.median([f or k + 1 for f in found]):.0f})") + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--index", required=True) + ap.add_argument("--synthetic", type=int, help="number of synthetic sketches") + ap.add_argument("--real", help="folder of real sketches named after their drawing") + ap.add_argument("--photo", action="store_true") + ap.add_argument("--model", help="evaluate this model instead of the index's (re-embeds renders in memory)") + ap.add_argument("--seed", type=int, default=1) + ap.add_argument("--save", help="folder to save synthetic sketches and their normalized form") + ap.add_argument("--aspect-weight", type=float, nargs="+", default=[ASPECT_WEIGHT], + help="one or more proportion-penalty weights to compare") + args = ap.parse_args() + k = 200 + index = SketchIndex(args.index, args.model) + + if args.synthetic: + rng = random.Random(args.seed) + holdout = [i for i, it in enumerate(index.items) if sc.is_holdout(it["key"])] + picks = rng.sample(holdout, min(args.synthetic, len(holdout))) + ranks = {w: [] for w in args.aspect_weight} + if args.save: + Path(args.save).mkdir(parents=True, exist_ok=True) + for n, i in enumerate(picks): + it = index.items[i] + sketch = sketchify(build_index.load_polys(args.index, it), rng) + mask = sc.sketch_mask(sketch) + for w in args.aspect_weight: + ranks[w].append(rank_of(index, full_ranking(index, mask, k, w), i)) + if args.save and n < 40: + sketch.save(Path(args.save, f"{n:03}-sketch.png")) + sc.normalize(mask).save(Path(args.save, f"{n:03}-query.png")) + Path(args.save, f"{n:03}-answer.txt").write_text( + json.dumps({"name": it["name"], "thumb": it["thumb"], "rank": ranks[args.aspect_weight[0]][-1]})) + for w in args.aspect_weight: + report(f"synthetic sketches vs {len(index.items)} drawings, aspect weight {w:g}", ranks[w], k) + + if args.real: + by_name = {} + for i, it in enumerate(index.items): + by_name.setdefault(it["name"].lower(), []).append(i) + ranks = [] + for f in sorted(Path(args.real).iterdir()): + if f.suffix.lower() not in (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tif", ".tiff"): + continue + name = f.stem.split("__")[0].lower() + if name not in by_name: + print(f" {f.name}: no indexed drawing named '{name}'") + continue + ranking = full_ranking(index, sc.sketch_mask(Image.open(f), photo=args.photo), k, args.aspect_weight[0]) + best = min((rank_of(index, ranking, i) for i in by_name[name]), + key=lambda r: (r[0] or k + 1, r[1] or k + 1)) + print(f" {f.name}: rank {best[0] or f'>{k}'} (same shape {best[1] or f'>{k}'})") + ranks.append(best) + if ranks: + report("real sketches", ranks, k) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/OpenNest.Training/sketch-search/requirements.txt b/OpenNest.Training/sketch-search/requirements.txt new file mode 100644 index 00000000..d363941a --- /dev/null +++ b/OpenNest.Training/sketch-search/requirements.txt @@ -0,0 +1,8 @@ +# PyTorch: install the CUDA build first for GPU embedding, e.g. +# pip install torch --index-url https://download.pytorch.org/whl/cu124 +torch>=2.4 +transformers>=4.45 +numpy>=1.26 +pillow>=10.0 +scikit-image>=0.22 +ezdxf>=1.3 diff --git a/OpenNest.Training/sketch-search/search.py b/OpenNest.Training/sketch-search/search.py new file mode 100644 index 00000000..ff441f32 --- /dev/null +++ b/OpenNest.Training/sketch-search/search.py @@ -0,0 +1,139 @@ +"""Finds the drawings that look most like a sketch. + + python search.py --index C:\\SketchIndex sketch.png + python search.py --index C:\\SketchIndex sketch.jpg --photo --size 10x4.5 --open + +Writes an HTML page of the best matches beside the index (or --out) and prints them. +""" +import argparse +import html +import json +import sys +import webbrowser +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +from PIL import Image + +import sketch_common as sc + +MM_PER_INCH = 25.4 +ASPECT_WEIGHT = 0.1 + + +class SketchIndex: + def __init__(self, folder, model=None): + """model overrides the index's model: its renders are re-embedded in memory.""" + self.folder = Path(folder) + config = json.loads(Path(folder, "config.json").read_text()) + self.items = [json.loads(l) for l in Path(folder, "items.jsonl").read_text(encoding="utf-8").splitlines()] + self.embedder = sc.Embedder(model or config["model"]) + if model and model != config["model"]: + load = lambda it: Image.open(Path(folder, "thumbs", it["thumb"])).convert("RGB") + with ThreadPoolExecutor(8) as pool: + self.embeddings = np.concatenate([self.embedder.embed(list(pool.map(load, self.items[i:i + 2048]))) + for i in range(0, len(self.items), 2048)]) + else: + self.embeddings = np.load(Path(folder, "embeddings.npy")).astype(np.float32) + self._aspects = None + + def size_mask(self, size, tolerance): + """Drawings whose extents match size (either orientation, inches or millimetres).""" + want = np.sort(np.array(size, dtype=np.float64)) + have = np.sort(np.array([[it["width"], it["height"]] for it in self.items]), axis=1) + ok = np.zeros(len(self.items), bool) + for factor in (1.0, MM_PER_INCH, 1 / MM_PER_INCH): + ok |= np.all(np.abs(have * factor - want) <= tolerance * want, axis=1) + return ok + + def scores(self, query, aspect=None, aspect_weight=ASPECT_WEIGHT): + """Similarity of every drawing to a normalized query image; aspect is the + sketch's short/long side ratio, which penalizes drawings of other proportions.""" + variants = self.embedder.embed(sc.dihedral(query)) + scores = (self.embeddings @ variants.T).max(axis=1) + if aspect and aspect_weight: + if self._aspects is None: + wh = np.array([[it["width"], it["height"]] for it in self.items]) + self._aspects = np.maximum(wh.min(axis=1), 1e-9) / np.maximum(wh.max(axis=1), 1e-9) + scores = scores - aspect_weight * np.abs(np.log(self._aspects / aspect)) + return scores + + def search(self, query, k=30, size=None, tolerance=0.15, aspect=None, aspect_weight=ASPECT_WEIGHT): + """query is a normalized image; returns [(score, item)] best first.""" + scores = self.scores(query, aspect, aspect_weight) + if size: + scores = np.where(self.size_mask(size, tolerance), scores, -np.inf) + order = np.argsort(-scores)[:k] + return [(float(scores[i]), self.items[i]) for i in order if np.isfinite(scores[i])] + + +def parse_size(text): + parts = text.lower().replace(" ", "").split("x") + if len(parts) != 2: + raise argparse.ArgumentTypeError("size looks like 10x4.5") + return float(parts[0]), float(parts[1]) + + +def write_html(path, index_folder, query_image, results, title): + path = Path(path) + query_file = path.with_suffix(".query.png") + query_image.save(query_file) + thumbs = Path(index_folder, "thumbs").resolve().as_uri() + cards = [] + for rank, (score, it) in enumerate(results, 1): + details = " · ".join(str(x) for x in ( + it.get("customer"), it.get("material"), + f"{it['thickness']:g} thk" if it.get("thickness") else None, + f"{it['width']:.3g} × {it['height']:.3g} {it.get('units', '')}".strip()) if x) + where = it.get("path") or f"server drawing {it.get('id')}" + cards.append(f"""
+{rank}. {html.escape(it['name'])}{html.escape(details)} +score {score:.3f} · {html.escape(where)}
""") + path.write_text(f"""{html.escape(title)} + +

{html.escape(title)}

+

The sketch as the model sees it (left), then the {len(results)} closest drawings.

+
{''.join(cards)}
""", encoding="utf-8") + return path + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("sketch", help="image of the sketch (PNG/JPG; crop off any window chrome)") + ap.add_argument("--index", required=True) + ap.add_argument("-k", type=int, default=30, help="number of matches") + ap.add_argument("--photo", action="store_true", help="camera photo or scan: adaptive thresholding") + ap.add_argument("--size", type=parse_size, help="overall size, e.g. 10x4.5 (either orientation, in or mm)") + ap.add_argument("--tolerance", type=float, default=0.15, help="relative size tolerance (default 0.15)") + ap.add_argument("--out", help="HTML results file (default: /last-search.html)") + ap.add_argument("--open", action="store_true", help="open the results in a browser") + args = ap.parse_args() + + mask = sc.sketch_mask(Image.open(args.sketch), photo=args.photo) + index = SketchIndex(args.index) + query = sc.normalize(mask) + results = index.search(query, args.k, args.size, args.tolerance, sc.aspect(mask)) + for rank, (score, it) in enumerate(results, 1): + print(f"{rank:3}. {score:.3f} {it['name']} [{it.get('customer', '')}] " + f"{it['width']:.3g}x{it['height']:.3g}") + out = write_html(args.out or Path(args.index, "last-search.html"), args.index, query, results, + f"Matches for {Path(args.sketch).name}") + print(f"results: {out}") + if args.open: + webbrowser.open(out.resolve().as_uri()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/OpenNest.Training/sketch-search/sketch_common.py b/OpenNest.Training/sketch-search/sketch_common.py new file mode 100644 index 00000000..e2d1eb0a --- /dev/null +++ b/OpenNest.Training/sketch-search/sketch_common.py @@ -0,0 +1,264 @@ +"""Shared geometry reading, image normalization and embedding for sketch search. + +Drawings and sketches go through the same normalize() so the model compares like +with like: ink is reduced to 1px centre lines, cropped, scaled to fit a square +canvas and re-stroked at a fixed width. +""" +import math +import re + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter +from skimage.filters import threshold_local +from skimage.morphology import binary_erosion, remove_small_objects, skeletonize + +CANVAS = 224 +MARGIN = 14 +STROKE = 3 +RENDER = 1024 +FILLED_RATIO = 0.35 + +# Motions with these tags are marks or open break cuts, not part outline. +SKIPPED_TAGS = {"SCRIBE", "ETCH", "BREAK", "LEADIN", "LEADOUT"} +_WORD = re.compile(r"([A-Z])\s*(-?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?)") + + +def is_holdout(key): + """About 10% of drawings, chosen by key, are never trained on and are what evaluation samples.""" + import hashlib + return hashlib.md5(key.encode()).digest()[0] < 26 + + +# ---------------------------------------------------------------- geometry + +def _arc_points(sx, sy, ex, ey, cx, cy, clockwise): + r = math.hypot(sx - cx, sy - cy) + a0 = math.atan2(sy - cy, sx - cx) + a1 = math.atan2(ey - cy, ex - cx) + sweep = a1 - a0 + if clockwise: + if sweep >= -1e-9: + sweep -= 2 * math.pi + elif sweep <= 1e-9: + sweep += 2 * math.pi + n = max(4, math.ceil(abs(sweep) / (math.pi / 36))) + return [(cx + r * math.cos(a0 + sweep * i / n), cy + r * math.sin(a0 + sweep * i / n)) + for i in range(1, n + 1)] + + +def parse_program(text): + """Returns the cut polylines of an OpenNest drawing program (G90/G91, G00-G03).""" + absolute = True + motion = 0 + x = y = 0.0 + polys = [] + current = None + for raw in text.splitlines(): + code, _, tag = raw.partition(":") + words = _WORD.findall(code.upper()) + if not words: + continue + values = {} + for letter, number in words: + if letter == "G": + g = int(float(number)) + if g == 90: + absolute = True + elif g == 91: + absolute = False + elif g in (0, 1, 2, 3): + motion = g + else: + values = None # G65 calls and other codes are not geometry + break + else: + values[letter] = float(number) + if not values or ("X" not in values and "Y" not in values): + continue + if absolute: + nx, ny = values.get("X", x), values.get("Y", y) + else: + nx, ny = x + values.get("X", 0.0), y + values.get("Y", 0.0) + skipped = tag.strip().upper().split(" ")[0] in SKIPPED_TAGS + if motion == 0 or skipped: + current = None + else: + if current is None: + current = [(x, y)] + polys.append(current) + if motion == 1: + current.append((nx, ny)) + else: + cx, cy = x + values.get("I", 0.0), y + values.get("J", 0.0) + current.extend(_arc_points(x, y, nx, ny, cx, cy, motion == 2)) + x, y = nx, ny + return [p for p in polys if len(p) > 1] + + +_DXF_SKIP_TYPES = {"TEXT", "MTEXT", "DIMENSION", "HATCH", "POINT", "ATTDEF", "ATTRIB", + "LEADER", "MULTILEADER", "IMAGE", "WIPEOUT", "VIEWPORT", "SOLID", "3DFACE"} + + +def read_dxf(path): + """Returns the polylines of a DXF's model space, skipping text, dimensions and marks.""" + import ezdxf + from ezdxf import path as ezpath + from ezdxf import recover + + doc, _ = recover.readfile(str(path)) + polys = [] + + def visit(entities, depth=0): + for e in entities: + kind = e.dxftype() + layer = (e.dxf.get("layer", "") or "").upper() + if kind in _DXF_SKIP_TYPES or "ETCH" in layer or "SCRIBE" in layer: + continue + if kind == "INSERT": + if depth < 8: + try: + visit(e.virtual_entities(), depth + 1) + except Exception: + pass + continue + try: + p = ezpath.make_path(e) + except (TypeError, ValueError, ezdxf.DXFError): + continue + pts = [(v.x, v.y) for v in p.flattening(0.002, segments=8)] + if len(pts) > 1: + polys.append(pts) + + visit(doc.modelspace()) + return polys + + +def bounds(polys): + pts = np.array([pt for p in polys for pt in p]) + lo, hi = pts.min(axis=0), pts.max(axis=0) + return lo, hi + + +# ---------------------------------------------------------------- images + +def render_mask(polys, size=RENDER, width=2): + """Rasterizes polylines to a boolean ink mask, aspect preserved.""" + lo, hi = bounds(polys) + span = max(hi[0] - lo[0], hi[1] - lo[1], 1e-9) + scale = (size - 8) / span + w = max(8, int((hi[0] - lo[0]) * scale) + 8) + h = max(8, int((hi[1] - lo[1]) * scale) + 8) + img = Image.new("L", (w, h), 0) + draw = ImageDraw.Draw(img) + for p in polys: + draw.line([((px - lo[0]) * scale + 4, h - 1 - ((py - lo[1]) * scale + 4)) for px, py in p], + fill=255, width=width) + return np.asarray(img) > 0 + + +def sketch_mask(image, photo=False): + """Converts a sketch image (scan, photo or MS Paint drawing) to a boolean ink mask.""" + if image.mode in ("RGBA", "LA") or "transparency" in image.info: + bg = Image.new("RGBA", image.size, (255, 255, 255, 255)) + image = Image.alpha_composite(bg, image.convert("RGBA")) + gray = np.asarray(image.convert("L"), dtype=np.float32) + if gray.max() - gray.min() < 10: + raise ValueError("The sketch image is blank.") + if photo: + gray = np.asarray(Image.fromarray(gray.astype(np.uint8)).filter(ImageFilter.GaussianBlur(1.5)), + dtype=np.float32) + block = max(15, (min(gray.shape) // 20) | 1) + ink = gray < threshold_local(gray, block, offset=10) + else: + # Anything clearly darker than the paper is ink. Otsu would split faint + # strokes from dark ones; the border is paper even when a fill dominates. + border = np.concatenate([gray[0], gray[-1], gray[:, 0], gray[:, -1]]) + paper = float(np.median(border)) + if paper < 128: # light lines on a dark background + gray, paper = 255 - gray, 255 - paper + darkest = float(np.percentile(gray, 0.1)) + ink = gray < paper - 0.35 * (paper - darkest) + return remove_small_objects(ink, max(16, ink.size // 200000)) + + +def aspect(mask): + """Short/long side ratio of the ink's bounding box.""" + ys, xs = np.nonzero(mask) + if len(xs) == 0: + raise ValueError("No ink found.") + h, w = ys.max() - ys.min() + 1, xs.max() - xs.min() + 1 + return min(h, w) / max(h, w) + + +def normalize(mask): + """Crops, centre-lines, scales and re-strokes an ink mask; returns a CANVAS x CANVAS RGB image.""" + ys, xs = np.nonzero(mask) + if len(xs) == 0: + raise ValueError("No ink found.") + mask = mask[ys.min():ys.max() + 1, xs.min():xs.max() + 1] + mask = np.pad(mask, 2) + if mask.mean() > FILLED_RATIO: # a filled silhouette: keep its boundary + mask = mask & ~binary_erosion(mask) + mask = skeletonize(mask) + h, w = mask.shape + s = (CANVAS - 2 * MARGIN) / max(h, w) + tw, th = max(1, round(w * s)), max(1, round(h * s)) + small = Image.fromarray(mask.astype(np.uint8) * 255).resize((tw, th), Image.BOX) + small = small.point(lambda v: 255 if v > 0 else 0).filter(ImageFilter.MaxFilter(STROKE)) + canvas = Image.new("L", (CANVAS, CANVAS), 0) + canvas.paste(small, ((CANVAS - tw) // 2, (CANVAS - th) // 2)) + return Image.eval(canvas, lambda v: 255 - v).convert("RGB") + + +def dihedral(image): + """The 8 rotations/mirrors of a square image.""" + out = [] + for flip in (False, True): + base = image.transpose(Image.FLIP_LEFT_RIGHT) if flip else image + for k in range(4): + out.append(base.rotate(90 * k)) + return out + + +# ---------------------------------------------------------------- embedding + +class Embedder: + def __init__(self, model_name="facebook/dinov2-small", device=None): + import torch + from transformers import AutoModel + + self.torch = torch + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.dtype = torch.float16 if self.device == "cuda" else torch.float32 + self.model = AutoModel.from_pretrained(model_name).to(self.device, self.dtype).eval() + self.model_name = model_name + + def embed(self, images, batch=64): + """L2-normalized embeddings (CLS token joined with mean patch token) as float32.""" + torch = self.torch + out = [] + with torch.inference_mode(): + for i in range(0, len(images), batch): + x = torch.from_numpy(np.stack([pixels(im) for im in images[i:i + batch]])) + out.append(features(self.model, x.to(self.device, self.dtype)).cpu().numpy()) + return np.concatenate(out) if out else np.zeros((0, 0), np.float32) + + +_MEAN = np.array([0.485, 0.456, 0.406], np.float32).reshape(3, 1, 1) +_STD = np.array([0.229, 0.224, 0.225], np.float32).reshape(3, 1, 1) + + +def pixels(image): + """A normalized RGB image as a CHW float32 array in the model's input scale.""" + return (np.asarray(image, dtype=np.float32).transpose(2, 0, 1) / 255 - _MEAN) / _STD + + +def features(model, x): + """Embeddings of a batch of pixels() tensors; differentiable for training.""" + import torch.nn.functional as F + import torch + + hidden = model(pixel_values=x).last_hidden_state + cls = F.normalize(hidden[:, 0].float(), dim=1) + patch = F.normalize(hidden[:, 1:].float().mean(dim=1), dim=1) + return F.normalize(torch.cat([cls, patch], dim=1), dim=1) diff --git a/OpenNest.Training/sketch-search/synthetic.py b/OpenNest.Training/sketch-search/synthetic.py new file mode 100644 index 00000000..ecb4ae3f --- /dev/null +++ b/OpenNest.Training/sketch-search/synthetic.py @@ -0,0 +1,69 @@ +"""Fake hand sketches generated from drawing geometry, for evaluation and training. + +Sketches get wobbly lines, up to 15% distorted proportions, about half their small +holes left out, a few degrees of tilt, random line weight and pen darkness, and +are sometimes filled in as a silhouette. +""" +import numpy as np +from PIL import Image, ImageDraw + +import sketch_common as sc + + +def sketchify(polys, rng, reorient=True): + """A fake hand sketch of a drawing as a grayscale image. With reorient=False the + sketch keeps the drawing's orientation apart from a few degrees of tilt.""" + lo, hi = sc.bounds(polys) + span = max(hi - lo) + kept = [] + for p in polys: + a = np.array(p) + closed = np.hypot(*(a[0] - a[-1])) < 1e-3 * span + if closed and max(a.max(0) - a.min(0)) < 0.08 * span and rng.random() < 0.5: + continue + kept.append(a) + if not kept: + kept = [np.array(p) for p in polys] + + angle = (rng.choice([0, 90, 180, 270]) if reorient else 0) + rng.gauss(0, 3) + t = np.radians(angle) + rot = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]]) + stretch = np.diag([rng.uniform(0.85, 1.15), 1.0]) + mirror = np.diag([-1.0, 1.0]) if reorient and rng.random() < 0.25 else np.eye(2) + m = rot @ stretch @ mirror + + nprng = np.random.default_rng(rng.randrange(1 << 30)) + amp = span * rng.uniform(0.001, 0.005) + shaky = [] + for a in kept: + seg = np.hypot(*np.diff(a, axis=0).T) + length = np.concatenate([[0], np.cumsum(seg)]) + n = max(8, int(length[-1] / (span * 0.004))) + s = np.linspace(0, length[-1], n) + pts = np.stack([np.interp(s, length, a[:, 0]), np.interp(s, length, a[:, 1])], axis=1) + noise = nprng.normal(size=(n + 30, 2)) + kernel = np.hanning(31) / np.hanning(31).sum() + smooth = np.stack([np.convolve(noise[:, i], kernel, "valid") for i in range(2)], axis=1)[:n] + smooth *= amp / (smooth.std() + 1e-12) + shaky.append((pts - lo) @ m.T + smooth) + + allpts = np.concatenate(shaky) + lo2, hi2 = allpts.min(0), allpts.max(0) + scale = 700 / max(hi2 - lo2) + w, h = int((hi2[0] - lo2[0]) * scale) + 100, int((hi2[1] - lo2[1]) * scale) + 100 + to_px = lambda a: [((x - lo2[0]) * scale + 50, h - ((y - lo2[1]) * scale + 50)) for x, y in a] + + if rng.random() < 0.2: # filled silhouette, holes by even-odd + ink = np.zeros((h, w), bool) + for a in shaky: + layer = Image.new("1", (w, h), 0) + ImageDraw.Draw(layer).polygon(to_px(a), fill=1) + ink ^= np.asarray(layer) + img = Image.fromarray(np.where(ink, 0, 255).astype(np.uint8)) + else: + img = Image.new("L", (w, h), 255) + draw = ImageDraw.Draw(img) + width, pen = rng.randint(2, 9), rng.randint(0, 120) + for a in shaky: + draw.line(to_px(a), fill=pen, width=width, joint="curve") + return img diff --git a/OpenNest.Training/sketch-search/train.py b/OpenNest.Training/sketch-search/train.py new file mode 100644 index 00000000..0874a97f --- /dev/null +++ b/OpenNest.Training/sketch-search/train.py @@ -0,0 +1,122 @@ +"""Fine-tunes the sketch-search embedding on fake sketches of indexed drawings. + +Each step pairs fresh synthetic sketches with their drawings' renderings and trains +the model to pick each sketch's own drawing out of the batch (symmetric InfoNCE). +Held-out drawings (sketch_common.is_holdout) and near-duplicate renderings, which +would be false negatives within a batch, are left out. No hand labelling is needed. + + python train.py --index C:\\SketchIndex --out C:\\SketchIndex\\model + python build_index.py --index C:\\SketchIndex --server http://barge.lan:8091 --model C:\\SketchIndex\\model +""" +import argparse +import json +import math +import random +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +import build_index +import sketch_common as sc +from synthetic import sketchify + +DUPLICATE = 0.985 + + +class SketchPairs(torch.utils.data.Dataset): + def __init__(self, index, items): + self.index, self.items = str(index), items + + def __len__(self): + return len(self.items) + + def __getitem__(self, i): + it = self.items[i] + clean = Image.open(Path(self.index, "thumbs", it["thumb"])).convert("RGB") + try: + polys = build_index.load_polys(self.index, it) + sketch = sc.normalize(sc.sketch_mask(sketchify(polys, random.Random(), reorient=False))) + except Exception: # noqa: BLE001 - a degenerate sketch falls back to the clean image + sketch = clean + return torch.from_numpy(sc.pixels(sketch)), torch.from_numpy(sc.pixels(clean)) + + +def distinct(items, embeddings, device): + """Indices of items whose rendering is not a near-duplicate of an earlier item's.""" + e = torch.from_numpy(embeddings).to(device, torch.float16) + keep = [] + for start in range(0, len(items), 4096): + block = e[start:start + 4096] @ e[:start + 4096].T + rows = torch.arange(block.shape[0], device=device) + start + earlier = torch.arange(block.shape[1], device=device)[None, :] < rows[:, None] + best = block.masked_fill(~earlier, -1).max(dim=1).values + keep += [start + i for i, dup in enumerate((best >= DUPLICATE).cpu().numpy()) if not dup] + return keep + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--index", required=True, help="index built with the base model") + ap.add_argument("--out", required=True, help="folder for the fine-tuned model") + ap.add_argument("--base", help="model to start from (default: the index's model)") + ap.add_argument("--epochs", type=int, default=5) + ap.add_argument("--batch", type=int, default=96) + ap.add_argument("--lr", type=float, default=2e-5) + ap.add_argument("--temperature", type=float, default=0.05) + ap.add_argument("--workers", type=int, default=12) + args = ap.parse_args() + + from transformers import AutoModel + + device = "cuda" if torch.cuda.is_available() else "cpu" + config = json.loads(Path(args.index, "config.json").read_text()) + items = [json.loads(l) for l in Path(args.index, "items.jsonl").read_text(encoding="utf-8").splitlines()] + embeddings = np.load(Path(args.index, "embeddings.npy")) + keep = distinct(items, embeddings, device) + train = [items[i] for i in keep if not sc.is_holdout(items[i]["key"])] + print(f"{len(items)} drawings, {len(keep)} distinct, training on {len(train)}") + + loader = torch.utils.data.DataLoader( + SketchPairs(args.index, train), batch_size=args.batch, shuffle=True, drop_last=True, + num_workers=args.workers, persistent_workers=args.workers > 0, pin_memory=True) + model = AutoModel.from_pretrained(args.base or config["model"]).to(device).train() + opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.05) + total = args.epochs * len(loader) + warmup = max(1, total // 20) + sched = torch.optim.lr_scheduler.LambdaLR( + opt, lambda s: min(1, (s + 1) / warmup) * 0.5 * (1 + math.cos(math.pi * min(1, s / total)))) + + step, started = 0, time.time() + for epoch in range(args.epochs): + losses, hits = [], [] + for sketches, cleans in loader: + x = torch.cat([sketches, cleans]).to(device, non_blocking=True) + with torch.autocast(device, dtype=torch.bfloat16): + f = sc.features(model, x) + fs, fc = f[:len(sketches)], f[len(sketches):] + logits = fs @ fc.T / args.temperature + target = torch.arange(len(fs), device=device) + loss = (F.cross_entropy(logits, target) + F.cross_entropy(logits.T, target)) / 2 + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + sched.step() + step += 1 + losses.append(loss.item()) + hits.append((logits.argmax(dim=1) == target).float().mean().item()) + if step % 50 == 0: + print(f" step {step}/{total} loss {np.mean(losses[-50:]):.3f} " + f"in-batch top1 {100 * np.mean(hits[-50:]):.1f}% ({time.time() - started:.0f}s)", flush=True) + print(f"epoch {epoch + 1}: loss {np.mean(losses):.3f}, in-batch top1 {100 * np.mean(hits):.1f}%", flush=True) + model.save_pretrained(args.out) + print(f"saved {args.out}") + + +if __name__ == "__main__": + sys.exit(main())