image-similarity/web/script.js

// Plumbing for the demos. All hashing happens in the worker (wasm).

const worker = new Worker('worker.js', { type: 'module' });
let nextId = 0;
const pending = new Map();

function rpc(op, args) {
	return new Promise((resolve, reject) => {
		const id = nextId++;
		pending.set(id, { resolve, reject });
		worker.postMessage({ id, op, args });
	});
}

worker.onmessage = (e) => {
	const { id, ok, result, error } = e.data;
	const promise = pending.get(id);
	pending.delete(id);
	if (!promise) return;
	if (ok) {
		promise.resolve(result);
	} else {
		promise.reject(new Error(error));
	}
};

// Helpers

function pngUrl(bytes) {
	return URL.createObjectURL(new Blob([bytes], { type: 'image/png' }));
}

function hamming(a, b) {
	let x = a ^ b;
	let count = 0;
	while (x) {
		x &= x - 1n;
		count++;
	}
	return count;
}

function hex64(hash) {
	return hash.toString(16).padStart(16, '0');
}

// Renders the top `bits` bits of a hash, msb first
function bitGrid(container, hash, bits = 64) {
	container.replaceChildren();
	for (let i = bits - 1; i >= 0; i--) {
		const cell = document.createElement('div');
		cell.className = (hash >> BigInt(i)) & 1n ? 'bit on' : 'bit off';
		container.appendChild(cell);
	}
}

// 8x8 grid of coefficient magnitudes, log scale, red negative / blue positive
function heatGrid(container, values) {
	container.replaceChildren();
	const max = Math.max(...values.map(Math.abs), 1e-9);
	for (const value of values) {
		const cell = document.createElement('div');
		const strength = Math.log1p(Math.abs(value)) / Math.log1p(max);
		const hue = value < 0 ? 4 : 215;
		cell.style.background = `hsl(${hue} 70% ${15 + strength * 45}%)`;
		cell.title = value.toFixed(1);
		container.appendChild(cell);
	}
}

function show(section) {
	document.querySelectorAll('.needs-image').forEach((el) => el.classList.add('visible'));
	if (section) document.getElementById(section).scrollIntoView({ behavior: 'smooth' });
}

// Demo state
let baseBuffer = null; // ArrayBuffer of the selected image
let baseUrl = null;
let baseHashes = null; // BigUint64Array [dct, median, phash]

// Pipeline demo

const imageContainers = document.getElementsByClassName('image-original');
const formImage = document.getElementById('dctimage');

function reset() {
	for (const container of imageContainers) {
		container.replaceChildren();
	}
	document.getElementById('image-resize').replaceChildren();
	document.getElementById('resize').classList.remove('resize');
}

async function loadImage(buffer, url) {
	reset();
	baseBuffer = buffer;
	baseUrl = url;
	for (const container of imageContainers) {
		const img = document.createElement('img');
		img.src = url;
		container.appendChild(img);
	}

	const data = await rpc('pipeline', { buffer });
	baseHashes = data.hashes;
	const [dct, median, phash] = data.hashes;

	// resize step
	const resized = document.createElement('img');
	resized.src = pngUrl(data.resize8);
	document.getElementById('image-resize').replaceChildren(resized);
	const original = document.getElementById('resize-original').children[0];
	if (original && original.width) {
		original.width = original.width;
		original.height = original.height;
	}

	// median
	bitGrid(document.getElementById('median-bits'), median);
	document.getElementById('median-hash').innerText = hex64(median);

	// dct
	heatGrid(document.getElementById('image-dct'), data.coefficients);
	bitGrid(document.getElementById('dct-sign-bits'), data.masks[0], 36);
	bitGrid(document.getElementById('dct-ordinal-bits'), data.masks[1], 28);
	document.getElementById('dct-hash').innerText = hex64(dct);

	// phash
	const small = document.createElement('img');
	small.src = pngUrl(data.resize32);
	small.className = 'pixelated med';
	document.getElementById('phash-resize').replaceChildren(small);
	heatGrid(document.getElementById('phash-lowfreq'), data.lowfreq);
	bitGrid(document.getElementById('phash-bits'), phash);
	document.getElementById('phash-hash').innerText = hex64(phash);

	show();
	resetMutation();
	runMutation();
}

formImage.addEventListener('change', () => {
	const file = formImage.files[0];
	if (!file) return;
	file.arrayBuffer().then((buffer) => loadImage(buffer, URL.createObjectURL(file)));
});

document.getElementById('resize-button').addEventListener('click', () => {
	document.getElementById('resize').classList.add('resize');
});

// Sample images, also used to seed the ranking pool
const samples = ['img/moon1.jpg', 'img/moon2.jpg', 'img/sunflower1.jpg', 'img/sunflower2.jpg'];
const sampleContainer = document.getElementById('sample-images');
for (const src of samples) {
	fetch(src)
		.then((response) => (response.ok ? response.arrayBuffer() : Promise.reject(response.status)))
		.then((buffer) => {
			const button = document.createElement('img');
			button.src = src;
			button.className = 'sample';
			button.addEventListener('click', () => loadImage(buffer, src));
			sampleContainer.appendChild(button);
			addToPool(src.split('/').pop(), buffer, src);
		})
		.catch(() => {});
}

// Mutation demo

const mutators = {
	flip: { label: 'Horizontal flip' },
	jpeg: { label: 'JPEG quality', min: 1, max: 100, step: 1, value: 50 },
	blur: { label: 'Gaussian blur', min: 0, max: 10, step: 0.1, value: 1.5 },
	sharpen: { label: 'Unsharp mask', min: 0, max: 10, step: 0.1, value: 1.5 },
	scale: { label: 'Rescale %', min: 5, max: 100, step: 1, value: 50 },
	crop: { label: 'Crop, keep %', min: 50, max: 100, step: 1, value: 90 },
	rotate: { label: 'Rotate deg', min: -45, max: 45, step: 0.5, value: 2 },
	hue: { label: 'Hue shift deg', min: 0, max: 180, step: 1, value: 30 },
	brightness: { label: 'Brightness', min: -100, max: 100, step: 1, value: 30 },
	contrast: { label: 'Contrast', min: -100, max: 100, step: 1, value: 25 },
	letterbox: { label: 'Letterbox bar %', min: 1, max: 40, step: 1, value: 10 },
	logo: { label: 'Logo size %', min: 1, max: 60, step: 1, value: 10 },
	noise: { label: 'Gaussian noise', min: 0, max: 50, step: 1, value: 10 },
};

const kindSelect = document.getElementById('mutator-kind');
const amountSlider = document.getElementById('mutator-amount');
const amountValue = document.getElementById('mutator-amount-value');

for (const [kind, config] of Object.entries(mutators)) {
	const option = document.createElement('option');
	option.value = kind;
	option.innerText = config.label;
	kindSelect.appendChild(option);
}

function resetMutation() {
	const config = mutators[kindSelect.value];
	if (config.min === undefined) {
		amountSlider.disabled = true;
		amountValue.innerText = '';
	} else {
		amountSlider.disabled = false;
		amountSlider.min = config.min;
		amountSlider.max = config.max;
		amountSlider.step = config.step;
		amountSlider.value = config.value;
		amountValue.innerText = config.value;
	}
}

let mutationTimer = null;
async function runMutation() {
	if (!baseBuffer) return;
	const kind = kindSelect.value;
	const amount = amountSlider.disabled ? 0 : Number(amountSlider.value);
	amountValue.innerText = amountSlider.disabled ? '' : amount;
	const data = await rpc('mutate', { buffer: baseBuffer, kind, amount });

	document.getElementById('compare-a').src = baseUrl;
	document.getElementById('compare-b').src = pngUrl(data.png);

	const names = ['dct', 'median', 'phash'];
	names.forEach((name, i) => {
		document.getElementById(`dist-${name}`).innerText = hamming(baseHashes[i], data.hashes[i]);
	});
}

function scheduleMutation() {
	clearTimeout(mutationTimer);
	mutationTimer = setTimeout(runMutation, 120);
}

kindSelect.addEventListener('change', () => {
	resetMutation();
	scheduleMutation();
});
amountSlider.addEventListener('input', scheduleMutation);
resetMutation();

document.getElementById('compare-slider').addEventListener('input', (e) => {
	document.getElementById('compare-b').style.clipPath = `inset(0 0 0 ${e.target.value}%)`;
});

// Ranking demo

const pool = []; // {name, url, hashes}
const poolContainer = document.getElementById('ranking-pool');
const resultsContainer = document.getElementById('ranking-results');
const algoSelect = document.getElementById('ranking-algo');
let queryIndex = null;

async function addToPool(name, buffer, url) {
	const { hashes } = await rpc('hashes', { buffer });
	const index = pool.length;
	pool.push({ name, url, hashes });

	const thumb = document.createElement('figure');
	const img = document.createElement('img');
	img.src = url;
	const caption = document.createElement('figcaption');
	caption.innerText = name;
	thumb.append(img, caption);
	thumb.addEventListener('click', () => {
		queryIndex = index;
		rank();
	});
	poolContainer.appendChild(thumb);
}

function rank() {
	if (queryIndex === null) return;
	const algo = Number(algoSelect.value);
	const query = pool[queryIndex];
	const ranked = pool
		.map((entry) => ({ entry, distance: hamming(query.hashes[algo], entry.hashes[algo]) }))
		.sort((a, b) => a.distance - b.distance)
		.slice(0, 10);

	resultsContainer.replaceChildren();
	for (const { entry, distance } of ranked) {
		const thumb = document.createElement('figure');
		const img = document.createElement('img');
		img.src = entry.url;
		const caption = document.createElement('figcaption');
		caption.innerText = `${distance}  ${entry.name}`;
		if (entry === query) thumb.className = 'query';
		thumb.append(img, caption);
		resultsContainer.appendChild(thumb);
	}
}

algoSelect.addEventListener('change', rank);

document.getElementById('ranking-files').addEventListener('change', (e) => {
	for (const file of e.target.files) {
		file.arrayBuffer().then((buffer) => addToPool(file.name, buffer, URL.createObjectURL(file)));
	}
});