`onboarding/sangetsu-lilycolor/scripts/enrich-pilot.py`

133 lines
#!/usr/bin/env python3
"""Pilot LOCAL swatch enrichment ($0) — reuses enrich-local-hybrid's Pillow palette +
qwen2.5vl on Mac2 ollama (127.0.0.1). Maps dominant hex -> interior-designer lexicon colorway.
Output: one JSONL row per swatch with SKU, hex[], bg_hex, colorway_name, styles[], material,
patterns[]. STAGES ONLY — no publish, no DB write, no Shopify. Reversible.

Usage: enrich-pilot.py <skus_file> <out_jsonl>
  skus_file: one mfr_sku per line (image lives at HENRY/<sku>.jpg)
"""
import sys, os, json, io, base64, time, subprocess, urllib.request

HENRY = "/Volumes/Henry/dw-lily-images"
OLLAMA = os.environ.get("ENRICH_OLLAMA_URL", "http://127.0.0.1:11434")
VL_MODEL = os.environ.get("ENRICH_VL_MODEL", "qwen2.5vl:7b")
HERE = os.path.dirname(os.path.abspath(__file__))
PALETTE_PY = os.path.expanduser("~/Projects/enrich-local-hybrid/enrich-palette.py")
LEXICON_PY = os.path.join(HERE, "lexicon-colorway.py")

def palette(img_path, k=6):
    r = subprocess.run([sys.executable, PALETTE_PY, img_path, str(k)],
                       capture_output=True, text=True, timeout=30)
    if r.returncode != 0 or not r.stdout:
        raise RuntimeError("palette failed: " + (r.stderr or "no output"))
    return json.loads(r.stdout)  # {palette:[{hex,percentage}], image_b64}

def lexicon_name(hex_str):
    r = subprocess.run([sys.executable, LEXICON_PY, hex_str],
                       capture_output=True, text=True, timeout=10)
    try:
        return json.loads(r.stdout).get("colorway")
    except Exception:
        return None

def vl(image_b64, pal):
    prompt = (
        "This wallcovering swatch image was pixel-sampled into these EXACT colors "
        f"(hex + area%): {json.dumps(pal)} . In the SAME ORDER give a designer color name "
        "for each hex. Also: backgroundIndex (0-based index of the base/background color), "
        "styles (interior-design styles), patterns (motif vocabulary like Damask/Floral/"
        "Geometric/Stripe/Grasscloth-texture), material (Grasscloth/Silk/Vinyl/Paper/"
        "Non-woven/etc), imageType (scan_swatch|scan_flatbed|photo_full|photo_crop|render), "
        "usable (false only if blank/corrupt/not a product), description (one sentence). "
        'Never use the word "Wallpaper" — say "Wallcovering".'
    )
    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"},
        "usable":{"type":"boolean"},"description":{"type":"string"}},
        "required":["colorNames","backgroundIndex","styles","patterns","material","usable"]}
    body = json.dumps({"model":VL_MODEL,"prompt":prompt,"images":[image_b64],
                       "stream":False,"format":schema,"keep_alive":"15m",
                       "options":{"temperature":0.1}}).encode()
    req = urllib.request.Request(OLLAMA + "/api/generate", data=body,
                                 headers={"Content-Type":"application/json"})
    with urllib.request.urlopen(req, timeout=300) as resp:
        raw = json.loads(resp.read().decode())
    return json.loads(raw["response"])

def clean(arr, cap=4):
    out, seen = [], set()
    for v in (arr or []):
        s = str(v or "").strip()
        if not s or s.lower() in ("none","n/a","na","null","undefined"): continue
        k = s.lower()
        if k in seen: continue
        seen.add(k); out.append(s)
        if len(out) >= cap: break
    return out

def enrich_one(sku):
    img = os.path.join(HENRY, sku + ".jpg")
    if not os.path.exists(img):
        return {"sku": sku, "error": "image_not_found"}
    t0 = time.time()
    samp = palette(img, 6)
    pal = samp.get("palette", [])
    res = vl(samp["image_b64"], pal)
    if res.get("usable") is False:
        return {"sku": sku, "usable": False, "wall_s": round(time.time()-t0,1)}
    names = res.get("colorNames", [])
    colors = [{"name": (names[i] if i < len(names) else "").strip(),
               "hex": c["hex"], "percentage": c.get("percentage")}
              for i, c in enumerate(pal)]
    bgi = res.get("backgroundIndex", 0)
    if not (isinstance(bgi, int) and 0 <= bgi < len(colors)): bgi = 0
    fg = [c for i, c in enumerate(colors) if i != bgi]
    dominant = (fg[0] if fg else (colors[0] if colors else {"hex": ""}))
    return {
        "sku": sku,
        "hex": [c["hex"] for c in colors],
        "bg_hex": colors[bgi]["hex"] if colors else "",
        "bg_name": colors[bgi]["name"] if colors else "",
        "dominant_hex": dominant["hex"],
        "colorway_name": lexicon_name(dominant["hex"]),
        "vl_color_names": [c["name"] for c in colors],
        "styles": clean(res.get("styles")),
        "material": (res.get("material") or "").strip(),
        "patterns": clean(res.get("patterns")),
        "imageType": res.get("imageType"),
        "description": (res.get("description") or "").strip().replace("Wallpaper","Wallcovering"),
        "_provider": "local-hybrid", "_cost": "$0 (local)",
        "wall_s": round(time.time()-t0, 1),
    }

def main():
    skus = [l.strip() for l in open(sys.argv[1]) if l.strip()]
    out_path = sys.argv[2]
    ok = fail = 0
    with open(out_path, "w") as fout:
        for i, sku in enumerate(skus, 1):
            try:
                row = enrich_one(sku)
                if row.get("error") or row.get("usable") is False:
                    fail += 1
                else:
                    ok += 1
                fout.write(json.dumps(row, ensure_ascii=False) + "\n")
                fout.flush()
                w = row.get("wall_s", "?")
                print(f"[{i}/{len(skus)}] {sku} ok={ok} fail={fail} {w}s "
                      f"cw={row.get('colorway_name')} mat={row.get('material')}", flush=True)
            except Exception as e:
                fail += 1
                fout.write(json.dumps({"sku": sku, "error": str(e)[:200]}) + "\n"); fout.flush()
                print(f"[{i}/{len(skus)}] {sku} ERROR {str(e)[:120]}", flush=True)
    print(f"\nDONE ok={ok} fail={fail} total={len(skus)}", flush=True)

if __name__ == "__main__":
    main()