Files

104 lines
2.9 KiB
JavaScript

import init, { hash_all, resize_preview, dct_coefficients, dct_masks, phash_lowfreq, mutate, suite_hashes, suite_labels, BkIndex } from './pkg/image_similarity.js';
const ready = init();
// The retrieval demos keep their trees in wasm memory between calls:
// a small one that gets visualized and the full experiment index.
let toyIndex = null;
let bigIndex = null;
// Average runtime of fn, sampled until enough wall time passed to
// outrun the clamped performance.now resolution
function time(fn) {
const start = performance.now();
let reps = 0;
do {
fn();
reps++;
} while (performance.now() - start < 20);
return (performance.now() - start) / reps;
}
onmessage = async (e) => {
const { id, op, args } = e.data;
await ready;
try {
let result;
switch (op) {
case 'pipeline': {
const bytes = new Uint8Array(args.buffer);
result = {
resize8: resize_preview(bytes, 8),
resize32: resize_preview(bytes, 32),
coefficients: dct_coefficients(bytes),
masks: dct_masks(bytes),
lowfreq: phash_lowfreq(bytes),
hashes: hash_all(bytes),
};
break;
}
case 'hashes':
result = { hashes: hash_all(new Uint8Array(args.buffer)) };
break;
case 'mutate': {
const png = mutate(new Uint8Array(args.buffer), args.kind, args.amount);
result = { png, hashes: hash_all(png) };
break;
}
case 'suite':
result = {
hashes: suite_hashes(new Uint8Array(args.buffer)),
labels: suite_labels().split('\n'),
};
break;
case 'toy_build': {
toyIndex?.free();
toyIndex = new BkIndex();
for (const key of args.keys) toyIndex.insert(key);
result = { structure: toyIndex.structure(), size: toyIndex.len() };
break;
}
case 'toy_query':
result = {
found: toyIndex.find(args.key, args.radius),
visited: toyIndex.trace(args.key, args.radius),
};
break;
case 'big_build': {
bigIndex?.free();
bigIndex = new BkIndex();
bigIndex.insert_bytes(new Uint8Array(args.buffer));
result = { size: bigIndex.len() };
break;
}
case 'big_query': {
const found = bigIndex.find(args.key, args.radius);
result = {
found,
compared: bigIndex.find_compared(args.key, args.radius),
treeMs: time(() => bigIndex.find(args.key, args.radius)),
scanMs: time(() => bigIndex.scan(args.key, args.radius)),
};
break;
}
// compared counts and found counts for every radius up to max,
// so the radius slider works without a round trip per step
case 'big_profile': {
const compared = [];
const found = [];
for (let t = 0; t <= args.max; t++) {
compared.push(bigIndex.find_compared(args.key, t));
found.push(bigIndex.find(args.key, t).length / 2);
}
result = { compared, found };
break;
}
default:
throw new Error(`unknown op ${op}`);
}
postMessage({ id, ok: true, result });
} catch (err) {
postMessage({ id, ok: false, error: String(err) });
}
};