← back to Designer Wallcoverings

gemini-classify-texture.js

113 lines

#!/usr/bin/env node
const { Client } = require('pg');
const https = require('https');
const http = require('http');

const GEMINI_KEY = '${GOOGLE_API_KEY}';
const GEMINI_URL = `https://generativelanguage.googleapis.com/v1beta/models/gemini-3.5-flash:generateContent?key=${GEMINI_KEY}`;
const PROMPT = `Is this wallcovering image a repeating pattern or a non-repeating texture? If it's a texture (grasscloth, linen, plain, solid, stucco, metallic, stone, concrete, weave, etc.), respond 'TEXTURE'. If it has a visible repeating pattern (florals, stripes, damask, geometric, medallion, toile, scenic, mural, etc.), respond 'PATTERN'. Respond with only one word: TEXTURE or PATTERN.`;

const db = new Client({ connectionString: (process.env.DATABASE_URL || 'postgresql://dw_admin@127.0.0.1:5432/dw_unified') });
const products = require('/tmp/thibaut_remaining.json');
const stats = { texture: 0, pattern: 0, error: 0, total: products.length };

function sleep(ms) { return new Promise(r => setTimeout(r, ms)); }

function downloadImage(url) {
  return new Promise((resolve, reject) => {
    const request = (targetUrl, redirectCount = 0) => {
      if (redirectCount > 5) return reject(new Error('Too many redirects'));
      const client = targetUrl.startsWith('https') ? https : http;
      client.get(targetUrl, { timeout: 15000 }, (res) => {
        if (res.statusCode >= 300 && res.statusCode < 400 && res.headers.location) {
          return request(res.headers.location, redirectCount + 1);
        }
        if (res.statusCode !== 200) { res.resume(); return reject(new Error(`HTTP ${res.statusCode}`)); }
        const chunks = [];
        res.on('data', chunk => chunks.push(chunk));
        res.on('end', () => {
          const buf = Buffer.concat(chunks);
          const ct = res.headers['content-type'] || 'image/jpeg';
          resolve({ base64: buf.toString('base64'), mimeType: ct.split(';')[0].trim() });
        });
        res.on('error', reject);
      }).on('error', reject).on('timeout', function() { this.destroy(); reject(new Error('timeout')); });
    };
    request(url);
  });
}

function classifyWithGemini(base64, mimeType) {
  const payload = {
    contents: [{ parts: [
      { text: PROMPT },
      { inline_data: { mime_type: mimeType, data: base64 } }
    ]}],
    generationConfig: { temperature: 0.1, maxOutputTokens: 10, thinkingConfig: { thinkingBudget: 0 } }
  };
  return new Promise((resolve, reject) => {
    const url = new URL(GEMINI_URL);
    const postData = JSON.stringify(payload);
    const req = https.request({
      hostname: url.hostname, path: url.pathname + url.search, method: 'POST',
      headers: { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(postData) },
      timeout: 30000
    }, (res) => {
      let data = '';
      res.on('data', chunk => data += chunk);
      res.on('end', () => {
        try {
          const json = JSON.parse(data);
          if (json.error) return reject(new Error(json.error.message));
          const text = json.candidates?.[0]?.content?.parts?.[0]?.text?.trim().toUpperCase() || '';
          if (text.includes('TEXTURE')) resolve('TEXTURE');
          else if (text.includes('PATTERN')) resolve('PATTERN');
          else resolve('UNKNOWN');
        } catch (e) { reject(e); }
      });
    });
    req.on('error', reject);
    req.on('timeout', () => { req.destroy(); reject(new Error('gemini timeout')); });
    req.write(postData); req.end();
  });
}

async function run() {
  await db.connect();
  console.log(`Classifying ${products.length} thibaut products with Gemini Vision (base64)...`);
  const startTime = Date.now();

  for (let i = 0; i < products.length; i++) {
    const p = products[i];
    try {
      const { base64, mimeType } = await downloadImage(p.image_url);
      const result = await classifyWithGemini(base64, mimeType);

      if (result === 'TEXTURE') {
        await db.query(`UPDATE thibaut_catalog SET repeat_h = '0', match_type = COALESCE(match_type, 'Texture'), updated_at = NOW() WHERE id = $1`, [p.id]);
        stats.texture++;
      } else if (result === 'PATTERN') {
        await db.query(`UPDATE thibaut_catalog SET match_type = COALESCE(match_type, 'Unknown Pattern'), updated_at = NOW() WHERE id = $1`, [p.id]);
        stats.pattern++;
      } else { stats.error++; }

      const pct = ((i + 1) / products.length * 100).toFixed(0);
      if ((i + 1) % 5 === 0 || i === products.length - 1)
        console.log(`  [${pct}%] ${i + 1}/${products.length} — T:${stats.texture} P:${stats.pattern} E:${stats.error} | ${p.pattern_name} → ${result}`);
      await sleep(200);
    } catch (err) {
      stats.error++;
      console.log(`  [ERR] #${p.id} ${p.pattern_name}: ${err.message.substring(0, 80)}`);
      await sleep(500);
    }
  }

  const elapsed = ((Date.now() - startTime) / 1000).toFixed(0);
  console.log(`\n=== GEMINI CLASSIFICATION COMPLETE (${elapsed}s) ===`);
  console.log(`Total: ${stats.total} | Texture: ${stats.texture} | Pattern: ${stats.pattern} | Errors: ${stats.error}`);
  await db.end();
  require('fs').writeFileSync('/tmp/gemini_classify_stats.json', JSON.stringify(stats));
}

run().catch(err => { console.error('Fatal:', err); process.exit(1); });