`enrich-local.js`

118 lines
// enrich-local.js — HYBRID local enrichment, drop-in companion to enrich-ai-tags.js.
// Ground-truth hex + percentages from real pixels (Pillow, more accurate than any VLM),
// qwen2.5vl for the semantic fields (color names, styles, patterns, material, imageType,
// dims, description, usability). Returns the SAME aiData shape as geminiAnalyze().
//
// Env:
//   ENRICH_OLLAMA_URL  Ollama base (default Mac1 via tailnet for Kamatera)
//   ENRICH_VL_MODEL    vision model (default qwen2.5vl:7b)
//   ENRICH_PY          python3 path (default 'python3')
const { spawnSync } = require('child_process');
const http = require('http');
const path = require('path');
const fs   = require('fs');
const os   = require('os');

const OLLAMA_URL = process.env.ENRICH_OLLAMA_URL || 'http://100.94.103.98:11434'; // Mac1 tailnet
const VL_MODEL   = process.env.ENRICH_VL_MODEL  || 'qwen2.5vl:7b';
const PY         = process.env.ENRICH_PY        || 'python3';
const PALETTE_PY = path.join(__dirname, 'enrich-palette.py');
const MAX_COLORS = 6;

function samplePalette(imgPath, k = MAX_COLORS) {
  const r = spawnSync(PY, [PALETTE_PY, imgPath, String(k)], { encoding: 'utf8', timeout: 25000, maxBuffer: 1 << 20 });
  if (r.status !== 0 || !r.stdout) throw new Error('palette failed: ' + (r.stderr || (r.error && r.error.message) || 'no output'));
  return JSON.parse(r.stdout); // [{hex, percentage}]
}

function ollamaVL(imageB64, palette) {
  const prompt =
    `This wallcovering/fabric image was pixel-sampled into these EXACT colors (hex + area%): ` +
    `${JSON.stringify(palette)} . In the SAME ORDER, give a designer color name for each hex. ` +
    `Also: backgroundIndex (0-based index of the base/background color in that list), styles, ` +
    `patterns, material, imageType (scan_swatch|scan_flatbed|photo_full|photo_crop|render), ` +
    `physicalWidthInches (number), physicalHeightInches (number), usable (false only if blank/` +
    `corrupt/not a product), rejectionReason (short, "" if usable), description (one sentence).`;
  const schema = { type: 'object', properties: {
    colorNames: { type: 'array', items: { type: 'string' } },
    backgroundIndex: { type: 'integer' },
    styles: { type: 'array', items: { type: 'string' } },
    patterns: { type: 'array', items: { type: 'string' } },
    material: { type: 'string' }, imageType: { type: 'string' },
    physicalWidthInches: { type: 'number' }, physicalHeightInches: { type: 'number' },
    usable: { type: 'boolean' }, rejectionReason: { type: 'string' }, description: { type: 'string' },
  }, required: ['colorNames','backgroundIndex','styles','patterns','material','imageType','physicalWidthInches','description','usable'] };
  // keep_alive holds the 7B model resident on Mac1 between calls so only the FIRST
  // request pays the cold-load (which over the Kamatera→Mac1 tailnet exceeded the old
  // 120s ceiling and timed out). 300s timeout absorbs that first cold load.
  const body = JSON.stringify({ model: VL_MODEL, prompt, images: [imageB64], stream: false, format: schema, keep_alive: '15m', options: { temperature: 0.1 } });
  return new Promise((resolve, reject) => {
    const u = new URL(OLLAMA_URL + '/api/generate');
    const req = http.request(
      { hostname: u.hostname, port: u.port || 11434, path: u.pathname, method: 'POST',
        headers: { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(body) }, timeout: 300000 },
      res => { let raw = ''; res.on('data', c => raw += c); res.on('end', () => {
        try { resolve(JSON.parse(JSON.parse(raw).response)); } catch (e) { reject(new Error('VL parse: ' + e.message)); }
      }); });
    req.on('error', reject);
    req.on('timeout', () => req.destroy(new Error('VL timeout')));
    req.write(body); req.end();
  });
}

// Drop junk values ("None"/"N/A"/empty), de-dup, cap to 4 — keeps tag sets clean.
function cleanList(arr) {
  if (!Array.isArray(arr)) return [];
  const seen = new Set(), out = [];
  for (const v of arr) {
    const s = String(v || '').trim();
    if (!s || /^(none|n\/?a|null|undefined)$/i.test(s)) continue;
    const key = s.toLowerCase();
    if (seen.has(key)) continue;
    seen.add(key); out.push(s);
    if (out.length >= 4) break;
  }
  return out;
}

// localAnalyze(imageUrlOrCandidates, mode, fetchBuffer) -> aiData (geminiAnalyze shape) | null
async function localAnalyze(imageUrl, mode, fetchBuffer) {
  const candidates = (Array.isArray(imageUrl) ? imageUrl : [imageUrl]).filter(Boolean);
  let buf = null;
  for (const c of candidates) { try { buf = await fetchBuffer(c); break; } catch (e) { /* next */ } }
  if (!buf) return null; // signal caller to fall back to Gemini
  const tmp = path.join(os.tmpdir(), `enrich-${process.pid}-${Date.now()}.img`);
  fs.writeFileSync(tmp, buf);
  try {
    const sampled = samplePalette(tmp, MAX_COLORS);          // {palette, image_b64 (normalized JPEG)}
    const palette = sampled.palette || [];
    if (!sampled.image_b64) throw new Error('palette: no normalized image');
    const vl = await ollamaVL(sampled.image_b64, palette);   // normalized → ollama always loads it
    if (vl.usable === false) {
      return { image_rejected: true, rejection_reason: vl.rejectionReason || 'not a usable product image' };
    }
    const names = Array.isArray(vl.colorNames) ? vl.colorNames : [];
    const colors = palette.map((c, i) => ({ name: (names[i] || '').trim(), hex: c.hex, percentage: c.percentage }));
    const bgIdx = (Number.isInteger(vl.backgroundIndex) && vl.backgroundIndex >= 0 && vl.backgroundIndex < colors.length) ? vl.backgroundIndex : 0;
    const fg = colors.filter((_, i) => i !== bgIdx);          // dominant = most prominent NON-background color
    const dominant = fg[0] || colors[0] || { hex: '' };
    return {
      backgroundColor: colors[bgIdx] ? colors[bgIdx].name : '',
      backgroundHex:   colors[bgIdx] ? colors[bgIdx].hex : '',
      dominantHex:     dominant.hex,
      colors,
      styles:   cleanList(vl.styles),
      patterns: cleanList(vl.patterns),
      material: vl.material || '',
      imageType: vl.imageType || null,
      physicalWidthInches:  vl.physicalWidthInches  || null,
      physicalHeightInches: vl.physicalHeightInches || null,
      description: vl.description || '',
      image_rejected: false,
      _provider: 'local-hybrid',
    };
  } finally { try { fs.unlinkSync(tmp); } catch (e) { /* ignore */ } }
}

module.exports = { localAnalyze, samplePalette };