mirror of
https://github.com/ajisaacs/OpenNest.git
synced 2026-10-28 16:01:06 -04:00
feat(training): add sketch-based drawing search prototype
Python tool that finds drawings resembling a hand sketch, photo or MS Paint outline. It indexes the OpenNest server's drawing library (read-only) and/or DXF folders, compares DINOv2 embeddings of normalized renderings across all 8 rotations/mirrors, and writes an HTML page of the closest matches with an optional overall-size filter. train.py fine-tunes the embedding contrastively on synthetic sketches of the library itself (no labelling); evaluate.py measures recall on held-out drawings or a folder of real sketches. On 40,757 server drawings, held-out synthetic sketches find the right shape in the top 20 75% of the time, up from 18% zero-shot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
ajandClaude Opus 5.5 committed
1 parent
fcb61f3b3e
commit
0be9e2201a
10 files changed
+1014
-1
No files matched your search
@@ -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.<Name>/` 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
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
__pycache__/
|
||||
@@ -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 `<index>/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 `<index>/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.
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
@@ -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
|
||||
@@ -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"""<figure><img src="{thumbs}/{it['thumb']}"><figcaption>
|
||||
<b>{rank}. {html.escape(it['name'])}</b><span>{html.escape(details)}</span>
|
||||
<small title="{html.escape(where)}">score {score:.3f} · {html.escape(where)}</small></figcaption></figure>""")
|
||||
path.write_text(f"""<!doctype html><meta charset="utf-8"><title>{html.escape(title)}</title>
|
||||
<style>
|
||||
:root {{ color-scheme: light dark; --bg:#fafafa; --card:#fff; --fg:#222; --muted:#666; --line:#ddd; }}
|
||||
@media (prefers-color-scheme: dark) {{ :root {{ --bg:#1b1b1b; --card:#262626; --fg:#eee; --muted:#aaa; --line:#3a3a3a; }} }}
|
||||
body {{ font: 14px system-ui, sans-serif; background: var(--bg); color: var(--fg); margin: 16px; }}
|
||||
header {{ display: flex; gap: 16px; align-items: center; margin-bottom: 16px; }}
|
||||
header img {{ width: 160px; border: 1px solid var(--line); background: #fff; }}
|
||||
main {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 12px; }}
|
||||
figure {{ margin: 0; background: var(--card); border: 1px solid var(--line); border-radius: 6px; padding: 8px; }}
|
||||
figure img {{ width: 100%; background: #fff; border-radius: 4px; }}
|
||||
figcaption {{ display: flex; flex-direction: column; gap: 2px; overflow-wrap: anywhere; }}
|
||||
span, small {{ color: var(--muted); }}
|
||||
</style>
|
||||
<header><img src="{query_file.name}"><div><h2>{html.escape(title)}</h2>
|
||||
<p>The sketch as the model sees it (left), then the {len(results)} closest drawings.</p></div></header>
|
||||
<main>{''.join(cards)}</main>""", 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: <index>/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())
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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())
|
||||
Reference in new issue
Block a user