@desert-ant-labs/shapes 0.1.0

package/dist/model.js

@@ -0,0 +1,342 @@
+import { fit } from "./fit.js";
+import { preprocess } from "./preprocess.js";
+export { outline } from "./shape.js";
+class Tensor {
+    shape;
+    floats;
+    packed;
+    palette;
+    length;
+    constructor(shape, floats, packed, palette) {
+        this.shape = shape;
+        this.floats = floats;
+        this.packed = packed;
+        this.palette = palette;
+        this.length = shape.reduce((a, b) => a * b, 1);
+    }
+    get(i) {
+        const f = this.floats;
+        if (f)
+            return f[i];
+        const byte = this.packed[i >> 1];
+        return this.palette[(i & 1) ? (byte >> 4) & 0xf : byte & 0xf];
+    }
+}
+// Minimal safetensors reader: u64 little-endian header length, JSON header, then
+// raw F32 tensors or packed 4-bit k-means palette tensors (`name` U8 indices +
+// `name.palette` F32 centroids, with logical shape in __metadata__).
+function parseSafetensors(bytes) {
+    const dv = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+    const headerLen = Number(dv.getBigUint64(0, true));
+    const header = JSON.parse(new TextDecoder().decode(bytes.subarray(8, 8 + headerLen)));
+    const dataStart = 8 + headerLen;
+    const meta = header.__metadata__ ?? {};
+    const out = new Map();
+    const slice = (name) => {
+        const e = header[name];
+        if (!e)
+            throw new Error(`shapes: missing tensor ${name}`);
+        const [a, b] = e.data_offsets;
+        const buf = bytes.slice(dataStart + a, dataStart + b);
+        return { shape: e.shape, bytes: buf };
+    };
+    const f32 = (name) => {
+        const s = slice(name);
+        return new Float32Array(s.bytes.buffer, s.bytes.byteOffset, s.bytes.byteLength / 4);
+    };
+    for (const name of Object.keys(header)) {
+        if (name === "__metadata__" || name.endsWith(".palette"))
+            continue;
+        if (header[name + ".palette"]) {
+            const shape = String(meta["shape." + name]).split(",").map(Number);
+            out.set(name, new Tensor(shape, null, slice(name).bytes, f32(name + ".palette")));
+        }
+        else {
+            const s = slice(name);
+            out.set(name, new Tensor(s.shape, new Float32Array(s.bytes.buffer, s.bytes.byteOffset, s.bytes.byteLength / 4), null, null));
+        }
+    }
+    return out;
+}
+function erf(x) {
+    const t = 1 / (1 + 0.3275911 * Math.abs(x));
+    const y = 1 -
+        (((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) *
+            t *
+            Math.exp(-x * x));
+    return x >= 0 ? y : -y;
+}
+const gelu = (x) => 0.5 * x * (1 + erf(x / Math.SQRT2));
+function layerNorm(v, w, b) {
+    const n = v.length;
+    let mean = 0;
+    for (let i = 0; i < n; i++)
+        mean += v[i];
+    mean /= n;
+    let varr = 0;
+    for (let i = 0; i < n; i++)
+        varr += (v[i] - mean) * (v[i] - mean);
+    varr /= n;
+    const inv = 1 / Math.sqrt(varr + 1e-5);
+    const out = new Float32Array(n);
+    for (let i = 0; i < n; i++)
+        out[i] = (v[i] - mean) * inv * w.get(i) + b.get(i);
+    return out;
+}
+/** Loads the portable model and runs stroke recognition. Construct once and reuse. */
+export class ShapesModel {
+    meta;
+    t;
+    D;
+    H;
+    dh;
+    scale;
+    layers;
+    constructor(weights, meta) {
+        this.meta = meta;
+        this.t = parseSafetensors(weights);
+        this.D = meta.model.width;
+        this.H = meta.model.heads;
+        this.dh = this.D / this.H;
+        this.scale = 1 / Math.sqrt(this.dh);
+        this.layers = meta.model.layers;
+    }
+    get(name) {
+        const t = this.t.get(name);
+        if (!t)
+            throw new Error(`shapes: missing tensor ${name}`);
+        return t;
+    }
+    /**
+     * Recognize a single stroke (ordered `[x, y]` points). Returns the snapped
+     * {@link Shape}, or `null` when the stroke is rejected or degenerate.
+     */
+    recognize(points) {
+        const feats = preprocess(points, this.meta.preprocess);
+        if (!feats)
+            return null;
+        const probs = this.classify(feats);
+        let best = 0;
+        for (let i = 1; i < probs.length; i++)
+            if (probs[i] > probs[best])
+                best = i;
+        const kind = this.meta.classes[best];
+        if (kind === "none")
+            return null;
+        const gate = this.meta.gates[kind];
+        if (gate && probs[best] < gate.conf)
+            return null;
+        const { shape, residual } = fit(kind, points);
+        if (gate && residual > gate.resid)
+            return null;
+        return shape;
+    }
+    // Conv1d stem -> Transformer encoder -> mean-pool -> MLP head -> softmax.
+    classify(feats) {
+        const N = feats.length;
+        // channel-major input [in_channels][N]
+        let h = [];
+        const cin = this.meta.model.in_channels;
+        for (let c = 0; c < cin; c++) {
+            const row = new Float32Array(N);
+            for (let t = 0; t < N; t++)
+                row[t] = feats[t][c];
+            h.push(row);
+        }
+        h = this.pool(this.relu(this.conv1d(h, "stem.conv1", cin, 32)));
+        h = this.pool(this.relu(this.conv1d(h, "stem.conv2", 32, 64)));
+        h = this.relu(this.conv1d(h, "stem.conv3", 64, 128));
+        const T = h[0].length;
+        const pe = this.get("encoder.pe");
+        // tokens [T][D]
+        const tok = [];
+        for (let t = 0; t < T; t++) {
+            const row = new Float32Array(this.D);
+            for (let d = 0; d < this.D; d++)
+                row[d] = h[d][t] + pe.get(t * this.D + d);
+            tok.push(row);
+        }
+        let layer = tok;
+        for (let l = 0; l < this.layers; l++)
+            layer = this.encoderLayer(layer, l);
+        const pooled = new Float32Array(this.D);
+        for (const row of layer)
+            for (let d = 0; d < this.D; d++)
+                pooled[d] += row[d];
+        for (let d = 0; d < this.D; d++)
+            pooled[d] /= layer.length;
+        const h0 = this.linear(pooled, "head.0", 64, this.D);
+        for (let i = 0; i < 64; i++)
+            h0[i] = Math.max(0, h0[i]);
+        const logits = this.linear(h0, "head.3", this.meta.classes.length, 64);
+        return softmax(logits);
+    }
+    conv1d(input, name, cin, cout) {
+        const W = this.get(name + ".weight");
+        const B = this.get(name + ".bias");
+        const T = input[0].length;
+        const out = [];
+        for (let o = 0; o < cout; o++) {
+            const row = new Float32Array(T);
+            const wbase = o * cin * 3;
+            for (let t = 0; t < T; t++) {
+                let acc = B.get(o);
+                for (let c = 0; c < cin; c++) {
+                    const ib = input[c];
+                    const wb = wbase + c * 3;
+                    if (t - 1 >= 0)
+                        acc += W.get(wb) * ib[t - 1];
+                    acc += W.get(wb + 1) * ib[t];
+                    if (t + 1 < T)
+                        acc += W.get(wb + 2) * ib[t + 1];
+                }
+                row[t] = acc;
+            }
+            out.push(row);
+        }
+        return out;
+    }
+    relu(x) {
+        for (const row of x)
+            for (let i = 0; i < row.length; i++)
+                if (row[i] < 0)
+                    row[i] = 0;
+        return x;
+    }
+    pool(x) {
+        const T2 = Math.floor(x[0].length / 2);
+        return x.map((row) => {
+            const out = new Float32Array(T2);
+            for (let i = 0; i < T2; i++)
+                out[i] = Math.max(row[2 * i], row[2 * i + 1]);
+            return out;
+        });
+    }
+    linear(x, name, out, inDim) {
+        const W = this.get(name + ".weight");
+        const B = this.get(name + ".bias");
+        const y = new Float32Array(out);
+        for (let o = 0; o < out; o++) {
+            let acc = B.get(o);
+            const base = o * inDim;
+            for (let i = 0; i < inDim; i++)
+                acc += W.get(base + i) * x[i];
+            y[o] = acc;
+        }
+        return y;
+    }
+    encoderLayer(tok, l) {
+        const p = `encoder.enc.layers.${l}.`;
+        const inW = this.get(p + "self_attn.in_proj_weight"); // [3D, D]
+        const inB = this.get(p + "self_attn.in_proj_bias");
+        const T = tok.length;
+        const D = this.D, H = this.H, dh = this.dh;
+        // q/k/v: [T][3D]
+        const q = [], k = [], v = [];
+        for (let t = 0; t < T; t++) {
+            const qkv = new Float32Array(3 * D);
+            for (let o = 0; o < 3 * D; o++) {
+                let acc = inB.get(o);
+                const base = o * D;
+                for (let i = 0; i < D; i++)
+                    acc += inW.get(base + i) * tok[t][i];
+                qkv[o] = acc;
+            }
+            q.push(qkv.subarray(0, D));
+            k.push(qkv.subarray(D, 2 * D));
+            v.push(qkv.subarray(2 * D, 3 * D));
+        }
+        // per-head scaled dot-product attention
+        const attn = [];
+        for (let t = 0; t < T; t++)
+            attn.push(new Float32Array(D));
+        for (let head = 0; head < H; head++) {
+            const off = head * dh;
+            for (let i = 0; i < T; i++) {
+                const scores = new Float32Array(T);
+                let mx = -Infinity;
+                for (let j = 0; j < T; j++) {
+                    let s = 0;
+                    for (let d = 0; d < dh; d++)
+                        s += q[i][off + d] * k[j][off + d];
+                    s *= this.scale;
+                    scores[j] = s;
+                    if (s > mx)
+                        mx = s;
+                }
+                let sum = 0;
+                for (let j = 0; j < T; j++) {
+                    scores[j] = Math.exp(scores[j] - mx);
+                    sum += scores[j];
+                }
+                const ai = attn[i];
+                for (let j = 0; j < T; j++) {
+                    const a = scores[j] / sum;
+                    for (let d = 0; d < dh; d++)
+                        ai[off + d] += a * v[j][off + d];
+                }
+            }
+        }
+        // out projection + residual + norm1
+        const outW = this.get(p + "self_attn.out_proj.weight");
+        const outB = this.get(p + "self_attn.out_proj.bias");
+        const n1w = this.get(p + "norm1.weight"), n1b = this.get(p + "norm1.bias");
+        const x1 = [];
+        for (let t = 0; t < T; t++) {
+            const a = new Float32Array(D);
+            for (let o = 0; o < D; o++) {
+                let acc = outB.get(o);
+                const base = o * D;
+                for (let i = 0; i < D; i++)
+                    acc += outW.get(base + i) * attn[t][i];
+                a[o] = acc + tok[t][o];
+            }
+            x1.push(layerNorm(a, n1w, n1b));
+        }
+        // FFN + residual + norm2
+        const l1w = this.get(p + "linear1.weight"), l1b = this.get(p + "linear1.bias");
+        const l2w = this.get(p + "linear2.weight"), l2b = this.get(p + "linear2.bias");
+        const ffn = l1b.length;
+        const n2w = this.get(p + "norm2.weight"), n2b = this.get(p + "norm2.bias");
+        const out = [];
+        for (let t = 0; t < T; t++) {
+            const hidden = new Float32Array(ffn);
+            for (let o = 0; o < ffn; o++) {
+                let acc = l1b.get(o);
+                const base = o * D;
+                for (let i = 0; i < D; i++)
+                    acc += l1w.get(base + i) * x1[t][i];
+                hidden[o] = gelu(acc);
+            }
+            const y = new Float32Array(D);
+            for (let o = 0; o < D; o++) {
+                let acc = l2b.get(o);
+                const base = o * ffn;
+                for (let i = 0; i < ffn; i++)
+                    acc += l2w.get(base + i) * hidden[i];
+                y[o] = acc + x1[t][o];
+            }
+            out.push(layerNorm(y, n2w, n2b));
+        }
+        return out;
+    }
+}
+function softmax(logits) {
+    let mx = -Infinity;
+    for (const x of logits)
+        if (x > mx)
+            mx = x;
+    let sum = 0;
+    const out = new Float32Array(logits.length);
+    for (let i = 0; i < logits.length; i++) {
+        out[i] = Math.exp(logits[i] - mx);
+        sum += out[i];
+    }
+    for (let i = 0; i < out.length; i++)
+        out[i] /= sum;
+    return out;
+}
+/** Creates a {@link ShapesModel} from raw buffers (lowest-level entry). */
+export function createShapes(buffers) {
+    return new ShapesModel(buffers.weights, buffers.meta);
+}

package/dist/preprocess.d.ts

@@ -0,0 +1,14 @@
+import type { Point } from "./shape.js";
+/** Frozen preprocessing constants (from the model's `shapes_meta.json`). */
+export interface PreprocessConfig {
+    spacing: number;
+    dist_mean: number;
+    dist_std: number;
+    add_curvature: boolean;
+}
+/**
+ * Port of the shared `preprocess.py`: dedupe → normalize → arc-length resample →
+ * `[dist, cos, sin]` features. Returns the `N × C` feature rows, or `null` for a
+ * degenerate stroke (too few points / too small).
+ */
+export declare function preprocess(points: Point[], cfg: PreprocessConfig): number[][] | null;

package/dist/preprocess.js

@@ -0,0 +1,111 @@
+const MIN_POINTS = 2;
+const MIN_TOTAL_LENGTH = 1e-6;
+const DEDUPE_EPSILON = 1e-9;
+/**
+ * Port of the shared `preprocess.py`: dedupe → normalize → arc-length resample →
+ * `[dist, cos, sin]` features. Returns the `N × C` feature rows, or `null` for a
+ * degenerate stroke (too few points / too small).
+ */
+export function preprocess(points, cfg) {
+    const cleaned = dedupe(points, DEDUPE_EPSILON);
+    if (cleaned.length < MIN_POINTS || totalLength(cleaned) < MIN_TOTAL_LENGTH)
+        return null;
+    const normalized = normalize(cleaned);
+    const resampled = resample(normalized, cfg.spacing);
+    return features(resampled, cfg.dist_mean, cfg.dist_std);
+}
+function dedupe(points, epsilon) {
+    if (points.length === 0)
+        return [];
+    const eps2 = epsilon * epsilon;
+    const out = [[points[0][0], points[0][1]]];
+    for (let i = 1; i < points.length; i++) {
+        const [px, py] = out[out.length - 1];
+        const dx = points[i][0] - px;
+        const dy = points[i][1] - py;
+        if (dx * dx + dy * dy > eps2)
+            out.push([points[i][0], points[i][1]]);
+    }
+    return out;
+}
+function totalLength(p) {
+    let total = 0;
+    for (let i = 1; i < p.length; i++)
+        total += Math.hypot(p[i][0] - p[i - 1][0], p[i][1] - p[i - 1][1]);
+    return total;
+}
+function normalize(p) {
+    let minX = p[0][0], maxX = p[0][0], minY = p[0][1], maxY = p[0][1];
+    for (const [x, y] of p) {
+        if (x < minX)
+            minX = x;
+        if (x > maxX)
+            maxX = x;
+        if (y < minY)
+            minY = y;
+        if (y > maxY)
+            maxY = y;
+    }
+    const cx = (minX + maxX) * 0.5;
+    const cy = (minY + maxY) * 0.5;
+    const longer = Math.max(maxX - minX, maxY - minY);
+    const scale = longer > 0 ? 1 / longer : 1;
+    return p.map(([x, y]) => [(x - cx) * scale, (y - cy) * scale]);
+}
+function resample(p, spacing) {
+    const n = p.length;
+    if (n <= 1)
+        return p.map(([x, y]) => [x, y]);
+    const out = [[p[0][0], p[0][1]]];
+    let [prevX, prevY] = p[0];
+    let since = 0;
+    let i = 1;
+    while (i < n) {
+        const dx = p[i][0] - prevX;
+        const dy = p[i][1] - prevY;
+        const segLen = Math.hypot(dx, dy);
+        if (segLen <= 0) {
+            prevX = p[i][0];
+            prevY = p[i][1];
+            i++;
+            continue;
+        }
+        const needed = spacing - since;
+        if (needed <= segLen) {
+            const t = needed / segLen;
+            prevX += dx * t;
+            prevY += dy * t;
+            out.push([prevX, prevY]);
+            since = 0;
+        }
+        else {
+            since += segLen;
+            prevX = p[i][0];
+            prevY = p[i][1];
+            i++;
+        }
+    }
+    const last = p[n - 1];
+    const [rx, ry] = out[out.length - 1];
+    if ((last[0] - rx) ** 2 + (last[1] - ry) ** 2 > 1e-18)
+        out.push([last[0], last[1]]);
+    return out;
+}
+function features(p, distMean, distStd) {
+    const std = distStd > 0 ? distStd : 1;
+    const out = [];
+    for (let i = 0; i < p.length; i++) {
+        let dist = 0, cos = 0, sin = 0;
+        if (i > 0) {
+            const dx = p[i][0] - p[i - 1][0];
+            const dy = p[i][1] - p[i - 1][1];
+            dist = Math.hypot(dx, dy);
+            if (dist > 0) {
+                cos = dx / dist;
+                sin = dy / dist;
+            }
+        }
+        out.push([(dist - distMean) / std, cos, sin]);
+    }
+    return out;
+}

package/dist/shape.d.ts

@@ -0,0 +1,34 @@
+/** A 2-D point, `[x, y]`, in the same coordinate space as the input stroke. */
+export type Point = [number, number];
+/** A recognized, fitted shape. */
+export type Shape = {
+    type: "line";
+    from: Point;
+    to: Point;
+} | {
+    type: "rectangle";
+    corners: Point[];
+} | {
+    type: "triangle";
+    vertices: Point[];
+} | {
+    type: "ellipse";
+    center: Point;
+    semiMajor: number;
+    semiMinor: number;
+    rotation: number;
+} | {
+    type: "star";
+    center: Point;
+    outerRadius: number;
+    innerRadius: number;
+    rotation: number;
+    pointCount: number;
+};
+/** The semantic kinds a stroke can be recognized as. */
+export type ShapeKind = Shape["type"];
+/**
+ * A closed (or, for a line, open) polyline outline of `shape`, suitable for
+ * rendering. `samples` controls the smoothness of the ellipse.
+ */
+export declare function outline(shape: Shape, samples?: number): Point[];

package/dist/shape.js

@@ -0,0 +1,38 @@
+/**
+ * A closed (or, for a line, open) polyline outline of `shape`, suitable for
+ * rendering. `samples` controls the smoothness of the ellipse.
+ */
+export function outline(shape, samples = 96) {
+    switch (shape.type) {
+        case "line":
+            return [shape.from, shape.to];
+        case "rectangle":
+            return shape.corners;
+        case "triangle":
+            return shape.vertices;
+        case "ellipse": {
+            const [cx, cy] = shape.center;
+            const c = Math.cos(shape.rotation);
+            const s = Math.sin(shape.rotation);
+            const out = [];
+            for (let i = 0; i < samples; i++) {
+                const t = (2 * Math.PI * i) / samples;
+                const x = shape.semiMajor * Math.cos(t);
+                const y = shape.semiMinor * Math.sin(t);
+                out.push([cx + x * c - y * s, cy + x * s + y * c]);
+            }
+            return out;
+        }
+        case "star": {
+            const [cx, cy] = shape.center;
+            const out = [];
+            const steps = shape.pointCount * 2;
+            for (let i = 0; i < steps; i++) {
+                const a = shape.rotation - Math.PI / 2 + (i * Math.PI) / shape.pointCount;
+                const r = i % 2 === 0 ? shape.outerRadius : shape.innerRadius;
+                out.push([cx + r * Math.cos(a), cy + r * Math.sin(a)]);
+            }
+            return out;
+        }
+    }
+}

package/package.json

@@ -0,0 +1,80 @@
+{
+  "name": "@desert-ant-labs/shapes",
+  "version": "0.1.0",
+  "description": "On-device single-stroke shape recognition.",
+  "type": "module",
+  "main": "./dist/index.node.js",
+  "types": "./dist/index.node.d.ts",
+  "exports": {
+    ".": {
+      "browser": {
+        "types": "./dist/index.browser.d.ts",
+        "default": "./dist/index.browser.js"
+      },
+      "worker": {
+        "types": "./dist/index.browser.d.ts",
+        "default": "./dist/index.browser.js"
+      },
+      "workerd": {
+        "types": "./dist/index.browser.d.ts",
+        "default": "./dist/index.browser.js"
+      },
+      "edge-light": {
+        "types": "./dist/index.browser.d.ts",
+        "default": "./dist/index.browser.js"
+      },
+      "node": {
+        "types": "./dist/index.node.d.ts",
+        "default": "./dist/index.node.js"
+      },
+      "default": {
+        "types": "./dist/index.node.d.ts",
+        "default": "./dist/index.node.js"
+      }
+    },
+    "./core": {
+      "types": "./dist/model.d.ts",
+      "default": "./dist/model.js"
+    },
+    "./package.json": "./package.json"
+  },
+  "files": [
+    "dist"
+  ],
+  "scripts": {
+    "build": "tsc",
+    "test": "node --import tsx --test test/*.test.ts",
+    "prepare": "npm run build",
+    "lint:publint": "publint --strict",
+    "lint:attw": "attw --pack . --profile esm-only",
+    "check:pkg": "npm run build && npm run lint:publint && npm run lint:attw"
+  },
+  "keywords": [
+    "shapes",
+    "shape-recognition",
+    "sketch",
+    "pen-input",
+    "pencilkit",
+    "on-device"
+  ],
+  "license": "SEE LICENSE IN LICENSE.md",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/Desert-Ant-Labs/shapes-js.git"
+  },
+  "homepage": "https://github.com/Desert-Ant-Labs/shapes-js",
+  "engines": {
+    "node": ">=18"
+  },
+  "sideEffects": false,
+  "publishConfig": {
+    "access": "public"
+  },
+  "devDependencies": {
+    "@arethetypeswrong/cli": "^0.18.4",
+    "@types/node": "^22",
+    "publint": "^0.3.21",
+    "tsx": "^4",
+    "typescript": "^5"
+  }
+}