@driftengine/capture 4.0.0

package/dist/models/twoWay.js

@@ -0,0 +1,67 @@
+/** The original SAM's names, which MobileSAM keeps. */
+export const SAM_NAMES = {
+    out: 'out_proj',
+    norm: (i) => `norm${i}`,
+    finalNorm: 'norm_final_attn',
+    mlp: ['lin1', 'lin2'],
+};
+/** Transformers' names for SAM 2. */
+export const SAM2_NAMES = {
+    out: 'o_proj',
+    norm: (i) => `layer_norm${i}`,
+    finalNorm: 'layer_norm_final_attn',
+    mlp: ['proj_in', 'proj_out'],
+};
+const EPSILON = 1e-5;
+/**
+ * The transformer under `prefix` over `tokens`, `[count, dim]`, and `image`, `[cells, dim]` in token
+ * layout with `positions` of the same shape; answers both, attended.
+ */
+export function twoWay(weights, graph, prefix, config, names, tokens, image, positions) {
+    const node = graph.node.bind(graph);
+    const { dim, heads } = config;
+    const read = weights.read.bind(weights);
+    const linear = (x, p, out, inputs) => node('linear', [x, read(`${p}weight`, [out, inputs]), read(`${p}bias`, [out])]);
+    const norm = (x, p) => node('layerNorm', [x, read(`${p}weight`, [dim]), read(`${p}bias`, [dim])], {
+        epsilon: EPSILON,
+    });
+    const attend = (p, q, k, v, narrow) => {
+        const inner = dim / narrow;
+        const attended = node('attention', [
+            linear(q, `${p}q_proj.`, inner, dim),
+            linear(k, `${p}k_proj.`, inner, dim),
+            linear(v, `${p}v_proj.`, inner, dim),
+        ], { heads });
+        return linear(attended, `${p}${names.out}.`, dim, inner);
+    };
+    const [into, outOf] = names.mlp;
+    let queries = tokens;
+    let keys = image;
+    for (let i = 0; i < config.depth; i += 1) {
+        const p = `${prefix}layers.${i}.`;
+        if (i === 0) {
+            queries = attend(`${p}self_attn.`, queries, queries, queries, 1);
+        }
+        else {
+            const q = node('add', [queries, tokens]);
+            queries = node('add', [queries, attend(`${p}self_attn.`, q, q, queries, 1)]);
+        }
+        queries = norm(queries, `${p}${names.norm(1)}.`);
+        let q = node('add', [queries, tokens]);
+        let k = node('add', [keys, positions]);
+        const toImage = attend(`${p}cross_attn_token_to_image.`, q, k, keys, config.downsample);
+        queries = norm(node('add', [queries, toImage]), `${p}${names.norm(2)}.`);
+        const hidden = node('relu', [linear(queries, `${p}mlp.${into}.`, config.mlp, dim)]);
+        const mixed = linear(hidden, `${p}mlp.${outOf}.`, dim, config.mlp);
+        queries = norm(node('add', [queries, mixed]), `${p}${names.norm(3)}.`);
+        q = node('add', [queries, tokens]);
+        k = node('add', [keys, positions]);
+        const toTokens = attend(`${p}cross_attn_image_to_token.`, k, q, queries, config.downsample);
+        keys = norm(node('add', [keys, toTokens]), `${p}${names.norm(4)}.`);
+    }
+    const q = node('add', [queries, tokens]);
+    const k = node('add', [keys, positions]);
+    const final = attend(`${prefix}final_attn_token_to_image.`, q, k, keys, config.downsample);
+    queries = norm(node('add', [queries, final]), `${prefix}${names.finalNorm}.`);
+    return { queries, keys };
+}

package/dist/models/vit.d.ts

@@ -0,0 +1,45 @@
+/**
+ * Depth Anything 3's backbone: DINOv2 as that model runs it, over one view or several at once.
+ *
+ * **Plain DINOv2 until `altStart`, and something else after it**, which is why this is not a
+ * generic vision transformer. From that block the class token of every view is replaced by a
+ * learned camera token — one for the first view, the reference, and another for the rest — and
+ * blocks alternate: an even one attends within each view, an odd one across all views at once.
+ * From `qkNormStart` each head's queries and keys are layer-normed, and from `ropeStart` they are
+ * turned by a two-dimensional rotary embedding: by the patch's row and column within a view, and
+ * across views by one position shared by every patch, so the global blocks see no geometry but
+ * which view a token belongs to.
+ *
+ * **Each tap is the last within-view output beside the current one**, `2·dim` wide, with only the
+ * second half normed — the upstream's `cat_token` — and its first token kept apart as that view's
+ * camera token, unnormed, which the camera decoder reads.
+ *
+ * **The positional embeddings are resized as the upstream resizes them**: bicubically, by the scale
+ * factor `(grid + 0.1)/trained`, whose inverse is the step a sample moves — not the ratio of the
+ * sizes. Only a square image at the trained grid skips it. The block norms take ε = 1e-6 and every
+ * other norm here PyTorch's default, 1e-5, which is a difference the checkpoint was trained with.
+ *
+ * What it gives up: the upstream reorders views around a chosen reference once there are three or
+ * more; here the first view given is the reference, and choosing it is the caller's.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+export interface VitConfig {
+    readonly dim: number;
+    readonly heads: number;
+    readonly depth: number;
+    readonly hidden: number;
+    readonly patch: number;
+    /** The patch grid the positional embeddings were trained at: 518 / 14 = 37. */
+    readonly trainedGrid: number;
+    readonly taps: readonly number[];
+    readonly altStart: number;
+    readonly qkNormStart: number;
+    readonly ropeStart: number;
+}
+export interface VitOutput {
+    /** Per tap, per view: the patch tokens, `[rows · cols, 2 · dim]`. */
+    readonly features: readonly (readonly string[])[];
+    /** Per tap: every view's camera token, `[views, 2 · dim]`. */
+    readonly cameras: readonly string[];
+}
+export declare function vitBackbone(weights: Weights, graph: GraphBuilder, config: VitConfig, prefix: string, images: readonly string[], rows: number, cols: number): VitOutput;

package/dist/models/vit.js

@@ -0,0 +1,157 @@
+import { applyRope, ropeConstants } from './rope.js';
+/** Block norms are trained with ε = 1e-6; the rest keep PyTorch's default. */
+const BLOCK_EPSILON = 1e-6;
+const DEFAULT_EPSILON = 1e-5;
+export function vitBackbone(weights, graph, config, prefix, images, rows, cols) {
+    const { dim, heads } = config;
+    const headDim = dim / heads;
+    const patches = rows * cols;
+    const tokens = patches + 1;
+    const w = (name, shape) => weights.read(`${prefix}${name}`, shape);
+    const node = graph.node.bind(graph);
+    /* Positions, resized to this grid unless it is the trained one on a square image. */
+    const trained = config.trainedGrid;
+    const table = node('reshape', [w('pos_embed', [1, trained * trained + 1, dim])], {
+        shape: [trained * trained + 1, dim],
+    });
+    let positions = table;
+    if (rows !== trained || cols !== trained) {
+        const grid = node('permute', [
+            node('reshape', [node('slice', [table], { axis: 0, start: 1, end: trained * trained + 1 })], {
+                shape: [trained, trained, dim],
+            }),
+        ], { order: [2, 0, 1] });
+        const resized = node('resize', [grid], {
+            height: rows,
+            width: cols,
+            mode: 'bicubic',
+            alignCorners: false,
+            stepHeight: trained / (rows + 0.1),
+            stepWidth: trained / (cols + 0.1),
+        });
+        const flat = node('reshape', [node('permute', [resized], { order: [1, 2, 0] })], {
+            shape: [patches, dim],
+        });
+        positions = node('concat', [node('slice', [table], { axis: 0, start: 0, end: 1 }), flat], {
+            axis: 0,
+        });
+    }
+    const cls = node('reshape', [w('cls_token', [1, 1, dim])], { shape: [1, dim] });
+    let views = images.map((image) => {
+        const embedded = node('patchEmbed', [
+            image,
+            w('patch_embed.proj.weight', [dim, 3, config.patch, config.patch]),
+            w('patch_embed.proj.bias', [dim]),
+        ], { patch: config.patch });
+        return node('add', [node('concat', [cls, embedded], { axis: 0 }), positions]);
+    });
+    /* Rotary tables: within a view by row and column from 1, across views all patches at 1, 1. */
+    const local = new Int32Array(2 * tokens);
+    const global = new Int32Array(2 * tokens * images.length);
+    for (let p = 0; p < patches; p += 1) {
+        local[2 * (p + 1)] = Math.floor(p / cols) + 1;
+        local[2 * (p + 1) + 1] = (p % cols) + 1;
+    }
+    for (let v = 0; v < images.length; v += 1) {
+        for (let p = 1; p < tokens; p += 1)
+            global.fill(1, 2 * (v * tokens + p), 2 * (v * tokens + p) + 2);
+    }
+    const rope = config.ropeStart >= 0
+        ? {
+            local: ropeConstants(weights, 'rope.local', local, heads, headDim),
+            global: ropeConstants(weights, 'rope.global', global, heads, headDim),
+        }
+        : null;
+    const block = (i, x, count, across) => {
+        const p = `blocks.${i}.`;
+        const h = node('layerNorm', [x, w(`${p}norm1.weight`, [dim]), w(`${p}norm1.bias`, [dim])], {
+            epsilon: BLOCK_EPSILON,
+        });
+        const qkv = node('linear', [
+            h,
+            w(`${p}attn.qkv.weight`, [3 * dim, dim]),
+            w(`${p}attn.qkv.bias`, [3 * dim]),
+        ]);
+        let q = node('slice', [qkv], { axis: 1, start: 0, end: dim });
+        let k = node('slice', [qkv], { axis: 1, start: dim, end: 2 * dim });
+        const v = node('slice', [qkv], { axis: 1, start: 2 * dim, end: 3 * dim });
+        if (config.qkNormStart >= 0 && i >= config.qkNormStart) {
+            const perHead = (value, norm) => node('reshape', [
+                node('layerNorm', [
+                    node('reshape', [value], { shape: [count * heads, headDim] }),
+                    w(`${p}attn.${norm}.weight`, [headDim]),
+                    w(`${p}attn.${norm}.bias`, [headDim]),
+                ], { epsilon: DEFAULT_EPSILON }),
+            ], { shape: [count, dim] });
+            q = perHead(q, 'q_norm');
+            k = perHead(k, 'k_norm');
+        }
+        if (rope !== null && i >= config.ropeStart) {
+            const tables = across ? rope.global : rope.local;
+            q = applyRope(graph, q, count, heads, headDim, tables.cos, tables.sin);
+            k = applyRope(graph, k, count, heads, headDim, tables.cos, tables.sin);
+        }
+        const attended = node('attention', [q, k, v], { heads });
+        const projected = node('linear', [
+            attended,
+            w(`${p}attn.proj.weight`, [dim, dim]),
+            w(`${p}attn.proj.bias`, [dim]),
+        ]);
+        const x1 = node('add', [x, node('mul', [projected, w(`${p}ls1.gamma`, [dim])])]);
+        const h2 = node('layerNorm', [x1, w(`${p}norm2.weight`, [dim]), w(`${p}norm2.bias`, [dim])], {
+            epsilon: BLOCK_EPSILON,
+        });
+        const inner = node('gelu', [
+            node('linear', [
+                h2,
+                w(`${p}mlp.fc1.weight`, [config.hidden, dim]),
+                w(`${p}mlp.fc1.bias`, [config.hidden]),
+            ]),
+        ]);
+        const out = node('linear', [
+            inner,
+            w(`${p}mlp.fc2.weight`, [dim, config.hidden]),
+            w(`${p}mlp.fc2.bias`, [dim]),
+        ]);
+        return node('add', [x1, node('mul', [out, w(`${p}ls2.gamma`, [dim])])]);
+    };
+    const features = [];
+    const cameras = [];
+    let latest = views;
+    for (let i = 0; i < config.depth; i += 1) {
+        if (config.altStart >= 0 && i === config.altStart) {
+            const camera = node('reshape', [w('camera_token', [1, 2, dim])], { shape: [2, dim] });
+            const reference = node('slice', [camera], { axis: 0, start: 0, end: 1 });
+            const source = node('slice', [camera], { axis: 0, start: 1, end: 2 });
+            views = views.map((x, v) => node('concat', [v === 0 ? reference : source, node('slice', [x], { axis: 0, start: 1, end: tokens })], { axis: 0 }));
+        }
+        const across = config.altStart >= 0 && i >= config.altStart && i % 2 === 1;
+        if (across) {
+            const all = views.length === 1 ? views[0] : node('concat', views, { axis: 0 });
+            const mixed = block(i, all, tokens * views.length, true);
+            views = views.map((_, v) => views.length === 1
+                ? mixed
+                : node('slice', [mixed], { axis: 0, start: v * tokens, end: (v + 1) * tokens }));
+        }
+        else {
+            views = views.map((x) => block(i, x, tokens, false));
+            latest = views;
+        }
+        if (config.taps.includes(i)) {
+            const tapped = views.map((x, v) => node('concat', [latest[v], x], { axis: 1 }));
+            cameras.push(tapped.length === 1
+                ? node('slice', [tapped[0]], { axis: 0, start: 0, end: 1 })
+                : node('concat', tapped.map((t) => node('slice', [t], { axis: 0, start: 0, end: 1 })), { axis: 0 }));
+            features.push(tapped.map((t) => {
+                const normed = node('layerNorm', [
+                    node('slice', [t], { axis: 1, start: dim, end: 2 * dim }),
+                    w('norm.weight', [dim]),
+                    w('norm.bias', [dim]),
+                ], { epsilon: DEFAULT_EPSILON });
+                const joined = node('concat', [node('slice', [t], { axis: 1, start: 0, end: dim }), normed], { axis: 1 });
+                return node('slice', [joined], { axis: 0, start: 1, end: tokens });
+            }));
+        }
+    }
+    return { features, cameras };
+}

package/dist/models/windowAttention.d.ts

@@ -0,0 +1,25 @@
+/**
+ * Attention within the windows of a grid, as TinyViT attends: each window a batch of tokens, each
+ * head's scores given a learned bias by the two tokens' distance.
+ *
+ * **Windows are partitioned as the upstream partitions them**: a grid the window does not divide is
+ * padded with zeros after its last row and column *before* the attention's own norm, so a padded
+ * token is the norm's β and is attended to as a key like any other — and the padding is sliced away
+ * afterwards. A grid exactly one window across is attended whole.
+ *
+ * **The bias table is laid out once, at conversion**, `[heads, n, n]` from the learned vector of
+ * `side²` distances, numbered in the order the upstream first meets them. What it gives up: the
+ * table costs MobileSAM's file 2.4 MB at half precision, where a gather on the device would keep the
+ * 3,000 learned values instead.
+ *
+ * **Queries, keys and values come out of one projection per head**, `[q k v]` for each head in turn
+ * where the attention operator takes all queries, then all keys, then all values; the projection's
+ * rows are reordered at conversion, which is where a head could otherwise be read as the wrong one's
+ * keys. The norm takes PyTorch's default ε, 1e-5.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+/**
+ * `[rows · cols, dim]` from `x` of that shape: the attention under `p` — `norm.`, `qkv.`, `proj.` and
+ * `attention_biases`, the upstream's names — within windows `side` tokens square.
+ */
+export declare function windowAttention(weights: Weights, graph: GraphBuilder, p: string, x: string, dim: number, heads: number, side: number, rows: number, cols: number): string;

package/dist/models/windowAttention.js

@@ -0,0 +1,104 @@
+const EPSILON = 1e-5;
+/**
+ * `[rows · cols, dim]` from `x` of that shape: the attention under `p` — `norm.`, `qkv.`, `proj.` and
+ * `attention_biases`, the upstream's names — within windows `side` tokens square.
+ */
+export function windowAttention(weights, graph, p, x, dim, heads, side, rows, cols) {
+    const node = graph.node.bind(graph);
+    const norm = (value, q) => node('layerNorm', [value, weights.read(`${q}weight`, [dim]), weights.read(`${q}bias`, [dim])], {
+        epsilon: EPSILON,
+    });
+    const linear = (value, q, out, inputs) => node('linear', [
+        value,
+        weights.read(`${q}weight`, [out, inputs]),
+        weights.read(`${q}bias`, [out]),
+    ]);
+    const n = side * side;
+    const whole = rows === side && cols === side;
+    const padRows = (side - (rows % side)) % side;
+    const padCols = (side - (cols % side)) % side;
+    const across = (rows + padRows) / side;
+    const down = (cols + padCols) / side;
+    const count = whole ? 1 : across * down;
+    let tokens = x;
+    if (!whole) {
+        let grid = node('reshape', [x], { shape: [rows, cols, dim] });
+        if (padRows > 0 || padCols > 0)
+            grid = node('pad', [grid], { after: [padRows, padCols, 0] });
+        const parts = node('reshape', [grid], { shape: [across, side, down, side, dim] });
+        tokens = node('reshape', [node('permute', [parts], { order: [0, 2, 1, 3, 4] })], {
+            shape: [count * n, dim],
+        });
+    }
+    const qkv = node('linear', [
+        norm(tokens, `${p}norm.`),
+        splitHeads(weights, `${p}qkv.weight`, [3 * dim, dim], heads),
+        splitHeads(weights, `${p}qkv.bias`, [3 * dim], heads),
+    ]);
+    const part = (i) => node('reshape', [node('slice', [qkv], { axis: 1, start: i * dim, end: (i + 1) * dim })], {
+        shape: [count, n, dim],
+    });
+    const attended = node('attention', [part(0), part(1), part(2), biasTable(weights, p, heads, side)], { heads });
+    const projected = linear(node('reshape', [attended], { shape: [count * n, dim] }), `${p}proj.`, dim, dim);
+    if (whole)
+        return projected;
+    const back = node('permute', [node('reshape', [projected], { shape: [across, down, side, side, dim] })], {
+        order: [0, 2, 1, 3, 4],
+    });
+    let grid = node('reshape', [back], { shape: [rows + padRows, cols + padCols, dim] });
+    if (padRows > 0)
+        grid = node('slice', [grid], { axis: 0, start: 0, end: rows });
+    if (padCols > 0)
+        grid = node('slice', [grid], { axis: 1, start: 0, end: cols });
+    return node('reshape', [grid], { shape: [rows * cols, dim] });
+}
+/**
+ * The projection `name` with its rows in the attention operator's order — every head's queries,
+ * then keys, then values — from the upstream's, where each head's three sit together.
+ */
+function splitHeads(weights, name, shape, heads) {
+    const rows = shape[0];
+    const width = rows / 3 / heads;
+    const columns = shape.length === 2 ? shape[1] : 1;
+    return weights.derive(name.replace(/\.(weight|bias)$/, '.split.$1'), shape, (values) => {
+        const from = values(name, shape);
+        const out = new Float32Array(from.length);
+        for (let head = 0; head < heads; head += 1) {
+            for (let part = 0; part < 3; part += 1) {
+                for (let j = 0; j < width; j += 1) {
+                    const source = (head * 3 + part) * width + j;
+                    const target = part * heads * width + head * width + j;
+                    out.set(from.subarray(source * columns, (source + 1) * columns), target * columns);
+                }
+            }
+        }
+        return out;
+    });
+}
+/**
+ * `[heads, n, n]`: each pair of a window's tokens given its head's learned bias for their distance
+ * in rows and in columns, the distances numbered in the order the upstream first meets them.
+ */
+function biasTable(weights, p, heads, side) {
+    const n = side * side;
+    return weights.derive(`${p}attention_bias_table`, [heads, n, n], (values) => {
+        const learned = values(`${p}attention_biases`, [heads, n]);
+        const numbered = new Map();
+        const index = new Int32Array(n * n);
+        for (let a = 0; a < n; a += 1) {
+            for (let b = 0; b < n; b += 1) {
+                const offset = Math.abs(Math.floor(a / side) - Math.floor(b / side)) * side +
+                    Math.abs((a % side) - (b % side));
+                if (!numbered.has(offset))
+                    numbered.set(offset, numbered.size);
+                index[a * n + b] = numbered.get(offset);
+            }
+        }
+        const out = new Float32Array(heads * n * n);
+        for (let head = 0; head < heads; head += 1) {
+            for (let i = 0; i < n * n; i += 1)
+                out[head * n * n + i] = learned[head * n + index[i]];
+        }
+        return out;
+    });
+}

package/dist/poses.d.ts

@@ -0,0 +1,31 @@
+export interface CaptureFrame {
+    readonly pixels: Uint8Array;
+    readonly width: number;
+    readonly height: number;
+}
+export interface PoseOptions {
+    /** `fx`, `fy`, `cx`, `cy` of the frames as they are given. */
+    readonly intrinsics: readonly [number, number, number, number];
+    /** The seeded generator RANSAC draws from; a capture answers the same twice. */
+    readonly random?: () => number;
+    /** Corners a frame. */
+    readonly budget?: number;
+    /** A depth model's own cameras, 3 × 4 world-to-camera each, where it ran. */
+    readonly modelPoses?: readonly ArrayLike<number>[] | null;
+    /** A distance in metres between two frames' places, which makes the scale metric. */
+    readonly knownLength?: {
+        readonly metres: number;
+        readonly from: number;
+        readonly to: number;
+    } | null;
+}
+export interface PoseResult {
+    readonly count: number;
+    /** Per frame, between zero and one: how much of it the answer rests on. */
+    readonly confidence: Float32Array;
+    /** False where the clip never moved: the rotations are answered and no baseline is claimed. */
+    readonly parallax: boolean;
+    readonly scale: 'metric' | 'relative';
+}
+/** The clip's camera path, into `out`: 3 × 4 world-to-camera a frame. */
+export declare function estimatePoses(frames: readonly CaptureFrame[], options: PoseOptions, out: Float32Array): PoseResult;