`scripts/enrich-colors.js`

173 lines
#!/usr/bin/env node
// TK-10112 — Classify each product's dominant color FROM ITS IMAGE so the Room
// Builder's 12-color filter is meaningful. products.color was title-derived, so
// almost everything landed in 'neutral' and the color swatches returned nothing.
//
// For each product (NOT is_wall_paint, image_url present): download the image,
// downsample to 16x16, average RGB across pixels while down-weighting near-white
// background pixels (product photos are usually shot on white), then map the
// average color -> HSV -> one of the 12-color-filter buckets and UPDATE
// products.color.
//
// Buckets (must match the Room Builder filter, server.js COLORS_12 + presets):
//   neutral, gray, black, brown, blue, green, pink, yellow, red
//
// Safe/reversible: only touches products.color. Resumable — pass --only-neutral
// to reprocess just the un-enriched rows. Concurrency-capped, 8s per-image
// timeout, skips (leaves color unchanged) on any download/decode failure.

const db = require('../lib/db');
const sharp = require('sharp');

const CONCURRENCY = 8;
const FETCH_TIMEOUT_MS = 8000;
const PROGRESS_EVERY = 200;

// ---- color math -----------------------------------------------------------

function rgbToHsv(r, g, b) {
  r /= 255; g /= 255; b /= 255;
  const max = Math.max(r, g, b), min = Math.min(r, g, b);
  const d = max - min;
  let h = 0;
  if (d !== 0) {
    if (max === r) h = ((g - b) / d) % 6;
    else if (max === g) h = (b - r) / d + 2;
    else h = (r - g) / d + 4;
    h *= 60;
    if (h < 0) h += 360;
  }
  const s = max === 0 ? 0 : d / max;
  const v = max;
  return { h, s, v };
}

// Map an average RGB to one of the 9 filter buckets.
function bucketFor(r, g, b) {
  const { h, s, v } = rgbToHsv(r, g, b);

  // Low-saturation → grayscale family (neutral / gray / black).
  if (s < 0.12) {
    if (v < 0.2) return 'black';
    if (v < 0.65) return 'gray';
    return 'neutral';
  }

  // Chromatic. A dark, warm, muted color reads as brown (wood/leather/taupe)
  // rather than a saturated red/orange.
  const warm = (h >= 345 || h < 45);
  if (v < 0.45 && warm) return 'brown';

  if (h >= 345 || h < 15) return 'red';       // red
  if (h < 45) return 'brown';                  // orange/amber → brown family
  if (h < 70) return 'yellow';                 // yellow
  if (h < 165) return 'green';                 // green
  if (h < 255) return 'blue';                  // cyan→blue
  if (h < 285) return 'blue';                  // purple → blue bucket (no purple swatch)
  if (h < 345) return 'pink';                  // magenta/pink
  return 'red';
}

// ---- per-image analysis ---------------------------------------------------

async function fetchImage(url) {
  const ctrl = new AbortController();
  const t = setTimeout(() => ctrl.abort(), FETCH_TIMEOUT_MS);
  try {
    const res = await fetch(url, {
      signal: ctrl.signal,
      headers: { 'User-Agent': 'Mozilla/5.0 (color-enrich/1.0)' },
    });
    if (!res.ok) return null;
    const ab = await res.arrayBuffer();
    return Buffer.from(ab);
  } catch {
    return null;
  } finally {
    clearTimeout(t);
  }
}

async function dominantBucket(buf) {
  // 16x16, alpha stripped, raw RGB. `fit: inside` preserves aspect so we don't
  // stretch. removeAlpha() flattens onto nothing — but raw() gives 3 channels.
  const { data, info } = await sharp(buf)
    .resize(16, 16, { fit: 'inside' })
    .removeAlpha()
    .raw()
    .toBuffer({ resolveWithObject: true });

  const ch = info.channels; // 3 after removeAlpha
  let rw = 0, gw = 0, bw = 0, wsum = 0;
  let rAll = 0, gAll = 0, bAll = 0, n = 0;

  for (let i = 0; i + ch - 1 < data.length; i += ch) {
    const r = data[i], g = data[i + 1], b = data[i + 2];
    rAll += r; gAll += g; bAll += b; n++;
    // Down-weight near-white pixels (photo background). Non-white subject
    // pixels get full weight so the SUBJECT color dominates the average.
    const nearWhite = r > 235 && g > 235 && b > 235;
    const w = nearWhite ? 0.08 : 1;
    rw += r * w; gw += g * w; bw += b * w; wsum += w;
  }
  if (n === 0) return null;

  // If the whole image is near-white (wsum collapsed to just background weight),
  // fall back to the plain average so we still classify it (as neutral/gray).
  let R, G, B;
  if (wsum < 0.5) {
    R = rAll / n; G = gAll / n; B = bAll / n;
  } else {
    R = rw / wsum; G = gw / wsum; B = bw / wsum;
  }
  return bucketFor(R, G, B);
}

// ---- driver ---------------------------------------------------------------

async function main() {
  const onlyNeutral = process.argv.includes('--only-neutral');
  const where = onlyNeutral
    ? "image_url IS NOT NULL AND NOT is_wall_paint AND (color IS NULL OR color = 'neutral')"
    : 'image_url IS NOT NULL AND NOT is_wall_paint';

  const { rows } = await db.query(
    `SELECT id, image_url FROM products WHERE ${where} ORDER BY id`
  );
  console.log(`[enrich-colors] ${rows.length} products to classify (concurrency ${CONCURRENCY})`);

  let done = 0, updated = 0, skipped = 0;
  const tally = {};

  let idx = 0;
  async function worker() {
    while (idx < rows.length) {
      const row = rows[idx++];
      try {
        const buf = await fetchImage(row.image_url);
        if (!buf) { skipped++; continue; }
        const bucket = await dominantBucket(buf);
        if (!bucket) { skipped++; continue; }
        await db.query('UPDATE products SET color = $1 WHERE id = $2', [bucket, row.id]);
        updated++;
        tally[bucket] = (tally[bucket] || 0) + 1;
      } catch {
        skipped++;
      } finally {
        done++;
        if (done % PROGRESS_EVERY === 0) {
          console.log(`  ${done}/${rows.length}  updated=${updated} skipped=${skipped}`);
        }
      }
    }
  }

  await Promise.all(Array.from({ length: CONCURRENCY }, worker));

  console.log(`[enrich-colors] done. updated=${updated} skipped=${skipped}`);
  console.log('[enrich-colors] this run bucketed:', JSON.stringify(tally));
  await db.pool.end();
}

main().catch((e) => { console.error(e); process.exit(1); });