artifex / js /engine.js
SuckSome's picture
Claude Opus 5.5
Deploy Artifex in-browser edition (model artifex-0.1-765e3b9a93, dataset 0.1+ac3dc1df9214)
a23acdf verified
Raw History Blame Contribute Delete
14.2 kB
// In-browser analysis engine, ported from src/artifex/predict.py, retrieve.py and uncertainty.py.
// Everything runs on the visitor's device; nothing about an upload is sent anywhere.
import * as store from "./store.js";
import { toModelArray } from "./preprocess.js";
import { randomId } from "./util.js";
let ort = null;
let sessionPromise = null;
async function loadOrt() {
if (!ort) {
ort = await import("../vendor/ort/ort.wasm.min.mjs");
ort.env.wasm.wasmPaths = new URL("../vendor/ort/", import.meta.url).href;
ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.max(1, Math.min(4, navigator.hardwareConcurrency || 1)) : 1;
}
return ort;
}
async function fetchModelBytes(url, expected, onProgress) {
const cacheName = "artifex-model";
let cache = null;
try {
cache = await caches.open(cacheName);
const hit = await cache.match(url);
if (hit) {
const buf = new Uint8Array(await hit.arrayBuffer());
if (buf.length === expected) {
onProgress?.(1, true);
return buf;
}
await cache.delete(url);
}
} catch {
cache = null; // Cache Storage unavailable (e.g. private mode): fall back to the HTTP cache
}
const r = await fetch(url);
if (!r.ok) throw new Error(`model download failed (HTTP ${r.status})`);
const buf = new Uint8Array(expected);
const reader = r.body.getReader();
let got = 0;
for (;;) {
const { done, value } = await reader.read();
if (done) break;
if (got + value.length > expected) throw new Error("model file is larger than expected");
buf.set(value, got);
got += value.length;
onProgress?.(got / expected, false);
}
if (got !== expected) throw new Error(`model download incomplete (${got} of ${expected} bytes)`);
if (cache) {
try {
// Old versions are removed so that only the current model occupies storage.
for (const req of await cache.keys()) if (req.url !== new URL(url, location.href).href) await cache.delete(req);
await cache.put(url, new Response(buf.slice(), { headers: { "content-type": "application/octet-stream" } }));
} catch {
/* storage quota: keep working without persistent caching */
}
}
return buf;
}
export function loadSession(onProgress) {
if (!sessionPromise) {
sessionPromise = (async () => {
const [m, o] = await Promise.all([store.model(), loadOrt()]);
const bytes = await fetchModelBytes(m.onnx_file, m.onnx_bytes, onProgress);
return o.InferenceSession.create(bytes, { executionProviders: ["wasm"], graphOptimizationLevel: "all" });
})().catch((e) => {
sessionPromise = null;
throw e;
});
}
return sessionPromise;
}
export async function embed(img) {
const [m, session] = await Promise.all([store.model(), loadSession()]);
const spec = m.encoder.transform;
const x = toModelArray(img, spec);
const out = await session.run({ pixel_values: new ort.Tensor("float32", x, [1, 3, spec.image_size, spec.image_size]) });
const q = out.embedding.data;
for (const v of q) if (!Number.isFinite(v)) throw new Error("encoder produced non-finite values");
return Float32Array.from(q);
}
// ---------------------------------------------------------------- classifier heads
function softmax(z, t) {
const s = z.map((v) => v / t);
const mx = Math.max(...s);
const e = s.map((v) => Math.exp(v - mx));
const sum = e.reduce((a, b) => a + b, 0);
return e.map((v) => v / sum);
}
export function headProba(head, q) {
const d = q.length;
const x = new Float64Array(d);
for (let i = 0; i < d; i++) x[i] = (q[i] - head.mean[i]) / head.scale[i];
const logits = head.coef.map((w, c) => {
let s = head.intercept[c];
for (let i = 0; i < d; i++) s += w[i] * x[i];
return s;
});
return softmax(logits, head.temperature);
}
function normalizedEntropy(p) {
if (p.length < 2) return 0;
let h = 0;
for (const v of p) {
const c = Math.min(1, Math.max(1e-12, v));
h -= c * Math.log(c);
}
return h / Math.log(p.length);
}
const NO_ABSTAIN_NOTE =
"Validation accuracy already exceeded the target without abstaining, so no confidence threshold applies to this head; " +
"accuracy is dominated by the largest classes, and small classes can still have low recall.";
const TARGET_MISSED_NOTE = "The validation accuracy target was not reachable for this head; the threshold enforces minimum coverage only.";
function headPrediction(name, head, q) {
const p = headProba(head, q);
const order = p.map((v, i) => i).sort((a, b) => p[b] - p[a] || a - b);
const ranked = order.map((i) => ({
label: head.classes[i],
display_name: head.display[i],
score: Math.min(1, Math.max(0, p[i])),
train_count: head.train_counts[head.classes[i]] || 0,
}));
const thr = head.abstention && head.abstention.threshold !== undefined ? head.abstention.threshold : null;
const top = p[order[0]];
const second = order.length > 1 ? p[order[1]] : 0;
let note = null;
if (head.abstention && head.abstention.target_met === false) note = TARGET_MISSED_NOTE;
else if (thr !== null && thr <= 0) note = NO_ABSTAIN_NOTE;
return {
pred: {
head: name, ranked, calibrated: true, calibration_method: "temperature scaling (validation split)",
abstained: thr !== null && top < thr, abstention_threshold: thr, top_score: top, margin: top - second, note,
},
p,
};
}
// ---------------------------------------------------------------- retrieval
export function similarities(E, q) {
const { data, n, dim } = E;
const out = new Float32Array(n);
for (let r = 0; r < n; r++) {
let s = 0;
const o = r * dim;
for (let i = 0; i < dim; i++) s += data[o + i] * q[i];
out[r] = s;
}
return out;
}
function rowDot(E, a, b) {
const { data, dim } = E;
let s = 0;
for (let i = 0; i < dim; i++) s += data[a * dim + i] * data[b * dim + i];
return s;
}
// ReferenceIndex.search: top-k, optional eligibility mask, near-identical results collapsed.
export function search(E, sims, k, { mask = null, collapse = null, exclude = null } = {}) {
const cand = [];
for (let i = 0; i < sims.length; i++) if ((!mask || mask[i]) && (exclude === null || i !== exclude)) cand.push(i);
cand.sort((a, b) => sims[b] - sims[a] || a - b);
const pre = cand.slice(0, Math.max(k * 8, 50));
const hits = [];
for (const i of pre) {
if (collapse !== null && hits.some((h) => rowDot(E, i, h.index) >= collapse)) continue;
hits.push({ index: i, similarity: sims[i] });
if (hits.length === k) break;
}
return hits;
}
function votes(E, sims, recs, column, k) {
const hits = search(E, sims, k);
const acc = new Map();
let total = 0;
for (const h of hits) {
const w = Math.max(h.similarity, 0);
const lab = recs[h.index][column];
const key = lab === null || lab === undefined ? "unlabelled" : String(lab);
acc.set(key, (acc.get(key) || 0) + w);
total += w;
}
if (total <= 0) return {};
const out = {};
for (const [key, v] of [...acc.entries()].sort((a, b) => b[1] - a[1])) out[key] = v / total;
return out;
}
export function similarRecord(rec, rank, sim, vocab) {
return {
artwork_id: rec.id, rank, similarity: Math.round(sim * 1e4) / 1e4, title: rec.title, artist_name: rec.artist,
date_display: rec.date, date_start: rec.ds, date_end: rec.de, region_normalized: rec.region,
region_display: rec.region ? vocab.region[rec.region] || rec.region : "Not recorded", culture_original: rec.culture_o,
period_label: rec.period, institution: rec.inst, source_url: rec.url, thumbnail_path: rec.thumb, image_license: rec.license,
};
}
export async function similar(q, sims, { region = null, period = null } = {}) {
const [m, E, R, vocab] = await Promise.all([store.model(), store.embeddings(), store.records(), store.vocab()]);
let mask = null;
if (region || period) mask = R.list.map((r) => (!region || r.region === region) && (!period || r.period === period));
return search(E, sims, m.retrieval.top_k, { mask, collapse: m.retrieval.collapse_cosine }).map((h, i) =>
similarRecord(R.list[h.index], i + 1, h.similarity, vocab),
);
}
// ---------------------------------------------------------------- analysis
const f2 = (x) => x.toFixed(2);
const f0 = (x) => x.toFixed(0);
export async function analyze(img, onStage) {
const t0 = performance.now();
const [m, E, R, vocab] = await Promise.all([store.model(), store.embeddings(), store.records(), store.vocab()]);
onStage?.("encode");
const q = await embed(img);
onStage?.("score");
const region = headPrediction("region", m.heads.region, q);
const period = headPrediction("period", m.heads.period, q);
let subregion = null;
const subHead = m.heads[`subregion:${region.pred.ranked[0].label}`];
if (subHead && !region.pred.abstained) subregion = headPrediction(`subregion:${region.pred.ranked[0].label}`, subHead, q).pred;
const sims = similarities(E, q);
const sim = await similar(q, sims);
const k = m.retrieval.vote_k;
const regionVotes = votes(E, sims, R.list, "region", k);
const periodVotes = votes(E, sims, R.list, "period", k);
let nearestRef = -Infinity;
let nearestTrain = -Infinity;
for (let i = 0; i < sims.length; i++) {
nearestRef = Math.max(nearestRef, sims[i]);
if (R.list[i].split === "train") nearestTrain = Math.max(nearestTrain, sims[i]);
}
if (nearestRef === -Infinity) nearestRef = 0;
if (nearestTrain === -Infinity) nearestTrain = 0;
const ood = m.ood ? nearestTrain < m.ood.threshold : false;
const unc = {
nearest_similarity: nearestRef, ood_threshold: m.ood ? m.ood.threshold : null, out_of_distribution: ood,
region_entropy: normalizedEntropy(region.p), period_entropy: normalizedEntropy(period.p),
};
const status = ood ? "out_of_scope" : region.pred.abstained && period.pred.abstained ? "insufficient_evidence" : "success";
const result = {
analysis_id: randomId(), status, model_version: m.model_version, dataset_version: m.dataset_version,
encoder: `${m.encoder.checkpoint}@${(m.encoder.commit || "").slice(0, 8)}`,
region: region.pred, subregion, period: period.pred, similar_artworks: sim,
retrieval_region_votes: regionVotes, retrieval_period_votes: periodVotes, uncertainty: unc,
evidence: evidence(m, vocab, region.pred, period.pred, subregion, regionVotes, periodVotes, unc),
image_width: img.width, image_height: img.height, elapsed_ms: Math.round((performance.now() - t0) * 10) / 10,
runtime: "browser (onnxruntime-web, WebAssembly)",
};
return { result, q, sims };
}
function displayName(vocab, kind, label) {
if (kind === "region") return vocab.region[label] || label;
if (kind === "period") return vocab.period[label] || label;
return vocab.subregion[label] || label;
}
function evidence(m, vocab, region, period, subregion, regionVotes, periodVotes, unc) {
const ev = [];
const k = m.retrieval.vote_k;
if (unc.out_of_distribution)
ev.push({ kind: "caution", text:
`The closest training image has similarity ${f2(unc.nearest_similarity)}, below the threshold ` +
`(${f2(unc.ood_threshold)}) exceeded by 99% of validation images. The upload may lie outside the ` +
"kinds of paintings the model was trained on; treat all scores below with particular caution." });
if (unc.nearest_similarity >= m.near_identical)
ev.push({ kind: "caution", text:
`A reference image is nearly identical to the upload (similarity ${f2(unc.nearest_similarity)}). The upload may be ` +
"a photograph of a work already in the reference collection; in that case the scores are not independent evidence." });
for (const [head, votesMap, noun] of [[region, regionVotes, "region"], [period, periodVotes, "period"]]) {
const top = head.ranked[0];
ev.push({ kind: "info", text:
`The ${noun} classifier ranks ${top.display_name} first with a calibrated score of ${f2(top.score)} ` +
`(runner-up ${head.ranked[1].display_name}, ${f2(head.ranked[1].score)}).` });
if (head.abstained)
ev.push({ kind: "caution", text:
`This ${noun} score is below the abstention threshold of ${f2(head.abstention_threshold)} chosen on validation data, ` +
`so Artifex reports insufficient evidence for the ${noun}.` });
const entries = Object.entries(votesMap);
if (entries.length) {
const [vtop, vshare] = entries[0];
const nMatch = Math.round((votesMap[top.label] || 0) * 100);
const vname = vtop === "unlabelled" ? "unlabelled works" : displayName(vocab, noun, vtop);
if (vtop === top.label)
ev.push({ kind: "agreement", text:
`Retrieval agrees: ${top.display_name} accounts for ${nMatch}% of the similarity-weighted votes ` +
`of the ${k} nearest reference works.` });
else
ev.push({ kind: "disagreement", text:
`Retrieval disagrees: the ${k} nearest reference works are mostly documented as ${vname} ` +
`(${Math.round(vshare * 100)}% of weighted votes), while ${top.display_name} accounts for ${nMatch}%.` });
}
if (top.train_count < m.limited_representation)
ev.push({ kind: "caution", text: `${top.display_name} has limited representation in the training data (${top.train_count} works).` });
const rec = (m.test_recall[noun] || {})[top.label];
if (rec)
ev.push({ kind: "info", text:
`On the held-out test split, the ${noun} classifier recovered ${f0(rec.recall * 100)}% of works labelled ` +
`${top.display_name} (n = ${rec.support}); precision when predicting it was ${f0(rec.precision * 100)}%.` });
}
if (subregion) {
const s = subregion.ranked[0];
ev.push({ kind: "info", text:
`Within ${region.ranked[0].display_name}, the tradition classifier ranks ${s.display_name} first (${f2(s.score)}).` +
(subregion.abstained ? " This score is below its abstention threshold." : "") });
}
const supported = m.heads.region.display.join(", ");
ev.push({ kind: "caution", text:
`Supported regions: ${supported}. Supported periods: ${m.heads.period.display.join(", ")}. Works from other traditions or periods ` +
"cannot be recognised and will be assigned to the nearest supported category." });
return ev;
}