storebound/cli/lib/imagehash.ts

// Perceptual image hashes without dependencies: a PNG decoder in TypeScript (node:zlib inflates), a 32 x 32 grey
// thumbnail, then dHash (gradients, 64 bit) and pHash (DCT low frequencies, 64 bit). JPEG / WebP go through ffmpeg
// when it is installed (one grey 32 x 32 frame on stdout), else they are skipped with a reason. Two images whose
// hashes differ in few bits look alike to a person (and to a reviewer comparing store pages).
import { readFileSync } from 'node:fs';
import { inflateSync } from 'node:zlib';
import * as rt from './rt';

export interface Gray { w: number; h: number; px: Float64Array }

/** Decodes a PNG to grey (alpha composited over white). 1-16 bit, every color type, Adam7 interlacing included. */
export function decodePng(buf: Buffer): Gray {
  if (buf.readUInt32BE(0) !== 0x89504e47) throw new Error('not a PNG');
  let w = 0, h = 0, depth = 0, ctype = 0, interlace = 0;
  let palette: Buffer | null = null, trns: Buffer | null = null;
  const idat: Buffer[] = [];
  for (let p = 8; p + 8 <= buf.length;) {
    const len = buf.readUInt32BE(p), type = buf.toString('latin1', p + 4, p + 8), data = buf.subarray(p + 8, p + 8 + len);
    if (type === 'IHDR') { w = data.readUInt32BE(0); h = data.readUInt32BE(4); depth = data[8]!; ctype = data[9]!; interlace = data[12]!; }
    else if (type === 'PLTE') palette = data;
    else if (type === 'tRNS') trns = data;
    else if (type === 'IDAT') idat.push(data);
    else if (type === 'IEND') break;
    p += 12 + len;
  }
  if (!w || !h) throw new Error('PNG without IHDR');
  if (w * h > 40_000_000) throw new Error('PNG too large to hash');
  const channels = { 0: 1, 2: 3, 3: 1, 4: 2, 6: 4 }[ctype];
  if (!channels) throw new Error(`PNG color type ${ctype} not supported`);
  const raw = inflateSync(Buffer.concat(idat));
  const bpp = Math.max(1, (channels * depth) >> 3);
  const out = new Float64Array(w * h);
  const max = (1 << depth) - 1;
  let off = 0;

  const pass = (x0: number, y0: number, dx: number, dy: number): void => {
    const pw = Math.ceil((w - x0) / dx), ph = Math.ceil((h - y0) / dy);
    if (pw <= 0 || ph <= 0) return;
    const stride = Math.ceil((pw * channels * depth) / 8);
    let prev = new Uint8Array(stride);
    for (let y = 0; y < ph; y++) {
      const ft = raw[off]!;
      const line = new Uint8Array(raw.subarray(off + 1, off + 1 + stride));
      off += 1 + stride;
      for (let i = 0; i < stride; i++) {
        const a = i >= bpp ? line[i - bpp]! : 0, b = prev[i]!, c = i >= bpp ? prev[i - bpp]! : 0;
        let v = line[i]!;
        if (ft === 1) v += a; else if (ft === 2) v += b; else if (ft === 3) v += (a + b) >> 1;
        else if (ft === 4) { const pp = a + b - c, pa = Math.abs(pp - a), pb = Math.abs(pp - b), pc = Math.abs(pp - c); v += pa <= pb && pa <= pc ? a : pb <= pc ? b : c; }
        line[i] = v & 255;
      }
      const sample = (x: number, ch: number): number => {
        if (depth === 8) return line[x * channels + ch]!;
        if (depth === 16) { const i = (x * channels + ch) * 2; return (line[i]! << 8) | line[i + 1]!; }
        const bit = (x * channels + ch) * depth;
        return (line[bit >> 3]! >> (8 - depth - (bit & 7))) & max;
      };
      for (let x = 0; x < pw; x++) {
        let r: number, g: number, bl: number, al = 1;
        if (ctype === 3) {
          const i = sample(x, 0);
          r = palette![i * 3]! / 255; g = palette![i * 3 + 1]! / 255; bl = palette![i * 3 + 2]! / 255;
          if (trns && i < trns.length) al = trns[i]! / 255;
        } else if (ctype === 0 || ctype === 4) { r = g = bl = sample(x, 0) / max; if (ctype === 4) al = sample(x, 1) / max; }
        else { r = sample(x, 0) / max; g = sample(x, 1) / max; bl = sample(x, 2) / max; if (ctype === 6) al = sample(x, 3) / max; }
        const lum = 0.299 * r + 0.587 * g + 0.114 * bl;
        out[(y0 + y * dy) * w + x0 + x * dx] = (lum * al + (1 - al)) * 255;
      }
      prev = line;
    }
  };
  if (interlace) for (const [x0, y0, dx, dy] of [[0, 0, 8, 8], [4, 0, 8, 8], [0, 4, 4, 8], [2, 0, 4, 4], [0, 2, 2, 4], [1, 0, 2, 2], [0, 1, 1, 2]] as const) pass(x0, y0, dx, dy);
  else pass(0, 0, 1, 1);
  return { w, h, px: out };
}

/** Area-average resize of a grey image. */
export function resize(g: Gray, W: number, H: number): Gray {
  const out = new Float64Array(W * H);
  for (let Y = 0; Y < H; Y++) for (let X = 0; X < W; X++) {
    const x0 = Math.floor((X * g.w) / W), x1 = Math.max(x0 + 1, Math.floor(((X + 1) * g.w) / W));
    const y0 = Math.floor((Y * g.h) / H), y1 = Math.max(y0 + 1, Math.floor(((Y + 1) * g.h) / H));
    let s = 0;
    for (let y = y0; y < y1; y++) for (let x = x0; x < x1; x++) s += g.px[y * g.w + x]!;
    out[Y * W + X] = s / ((x1 - x0) * (y1 - y0));
  }
  return { w: W, h: H, px: out };
}

/** 64-bit hashes as 16 hex characters. */
export interface ImageHash { d: string; p: string }

const toHex = (bits: number[]): string => { let s = ''; for (let i = 0; i < 64; i += 4) s += ((bits[i]! << 3) | (bits[i + 1]! << 2) | (bits[i + 2]! << 1) | bits[i + 3]!).toString(16); return s; };

export function dHash(g: Gray): string {
  const s = resize(g, 9, 8);
  const bits: number[] = [];
  for (let y = 0; y < 8; y++) for (let x = 0; x < 8; x++) bits.push(s.px[y * 9 + x]! < s.px[y * 9 + x + 1]! ? 1 : 0);
  return toHex(bits);
}

export function pHash(g: Gray): string {
  const N = 32, s = resize(g, N, N);
  const cos = Array.from({ length: 8 }, (_, u) => Array.from({ length: N }, (_, x) => Math.cos(((2 * x + 1) * u * Math.PI) / (2 * N))));
  const coef: number[] = [];
  for (let u = 0; u < 8; u++) for (let v = 0; v < 8; v++) {
    let sum = 0;
    for (let y = 0; y < N; y++) for (let x = 0; x < N; x++) sum += s.px[y * N + x]! * cos[u]![y]! * cos[v]![x]!;
    coef.push(sum);
  }
  const ac = coef.slice(1).sort((a, b) => a - b);
  const median = ac[Math.floor(ac.length / 2)]!;
  return toHex(coef.map((c, i) => (i === 0 ? 0 : c > median ? 1 : 0)));
}

export const hashGray = (g: Gray): ImageHash => ({ d: dHash(g), p: pHash(g) });

export function hamming(a: string, b: string): number {
  let n = 0;
  for (let i = 0; i < a.length; i++) { let x = parseInt(a[i]!, 16) ^ parseInt(b[i]!, 16); while (x) { n += x & 1; x >>= 1; } }
  return n;
}

/** 1 = same picture, about 0.5 = unrelated. The mean of both hashes. */
export const imageSimilarity = (a: ImageHash, b: ImageHash): number => 1 - (hamming(a.d, b.d) + hamming(a.p, b.p)) / 128;

/** Hash of an image file: PNG in TypeScript, anything else through ffmpeg (null with a reason when that is missing). */
export function hashFile(path: string): { hash: ImageHash | null; reason?: string } {
  try {
    const buf = readFileSync(path);
    if (buf.length > 8 && buf.readUInt32BE(0) === 0x89504e47) return { hash: hashGray(decodePng(buf)) };
    const ffmpeg = rt.which('ffmpeg');
    if (!ffmpeg) return { hash: null, reason: 'not a PNG and ffmpeg is not installed' };
    const r = rt.spawnSync([ffmpeg, '-v', 'error', '-i', path, '-frames:v', '1', '-vf', 'scale=32:32:flags=area,format=gray', '-f', 'rawvideo', '-'], { stdout: 'pipe', stderr: 'pipe', timeout: 20_000 });
    const px = r.stdout;
    if (r.exitCode !== 0 || px.length !== 1024) return { hash: null, reason: `ffmpeg could not read it (${r.stderr.toString().trim().split('\n')[0]?.slice(0, 120) ?? r.exitCode})` };
    return { hash: hashGray({ w: 32, h: 32, px: Float64Array.from(px) }) };
  } catch (e) { return { hash: null, reason: (e as Error).message }; }
}