`enrich.py`

109 lines
#!/usr/bin/env python3
"""
Japan microsite import enricher — per Steve's standing rule (2026-06-30):
extract hex# / style / material for every item so the catalog is searchable.

Engine: LOCAL HYBRID ($0) — Pillow ground-truth hex + qwen2.5vl semantics.
  hex     : Pillow median-cut palette (dominant colorway colors + coverage %)
  style   : qwen2.5vl, constrained to the interior-designer style lexicon
  material: qwen2.5vl, constrained to the surface-material lexicon

Usage: enrich.py <images_dir> <staging.jsonl> <out.jsonl> [--limit N] [--ollama URL]
Maps an image to a record by the SKU token in the filename (LIS-42003 -> LIS42003).
"""
import sys, os, json, ast, base64, re, glob, urllib.request, colorsys

OLLAMA = "http://127.0.0.1:11434"   # Mac2 failover default (Mac1 LAN wedged)
VMODEL = "qwen2.5vl:7b"

# interior-designer lexicon (constrained vocab so facets stay clean/searchable)
STYLES = ["Traditional","Modern","Floral","Damask","Geometric","Stripe","Texture/Plain",
          "Chinoiserie","Toile","Ikat","Botanical","Abstract","Animal","Trellis","Moiré",
          "Grasscloth/Natural","Metallic","Childrens","Mural/Scenic"]
MATERIALS = ["Vinyl","Non-woven","Paper","Grasscloth","Silk","Linen/Fabric-backed",
             "Cork","Mica/Mineral","Metallic foil","Wood veneer","Leather","Glass bead"]

def norm_sku(s): return re.sub(r'[^A-Za-z0-9]','',(s or '')).upper()

def hex_palette(path, k=6):
    from PIL import Image
    im = Image.open(path).convert("RGB")
    im.thumbnail((200,200))
    q = im.quantize(colors=k, method=Image.MEDIANCUT)
    pal = q.getpalette(); counts = q.getcolors() or []
    total = sum(c for c,_ in counts) or 1
    out=[]
    for cnt, idx in sorted(counts, reverse=True):
        r,g,b = pal[idx*3:idx*3+3]
        out.append({"hex":"#%02X%02X%02X"%(r,g,b), "pct":round(100*cnt/total,1)})
    return out

def dominant_family(hexes):
    """coarse color-family bucket for the nav facet, from the top swatch"""
    if not hexes: return None
    h=hexes[0]["hex"].lstrip("#")
    r,g,b=[int(h[i:i+2],16)/255 for i in (0,2,4)]
    hh,ss,vv=colorsys.rgb_to_hsv(r,g,b)
    if vv<0.18: return "Black"
    if ss<0.12 and vv>0.85: return "White"
    if ss<0.15: return "Grey"
    if ss<0.30 and vv>0.55: return "Neutral/Beige"
    deg=hh*360
    for lo,hi,name in [(0,15,"Red"),(15,45,"Orange/Brown"),(45,70,"Yellow"),
                       (70,160,"Green"),(160,200,"Teal"),(200,255,"Blue"),
                       (255,290,"Purple"),(290,335,"Pink"),(335,360,"Red")]:
        if lo<=deg<hi: return name
    return "Red"

def vision_tags(path):
    with open(path,"rb") as f: b64=base64.b64encode(f.read()).decode()
    prompt=(f"This is a wallcovering product swatch. Choose the single best STYLE from "
            f"{STYLES} and the single best MATERIAL from {MATERIALS}. "
            f'Reply ONLY compact JSON: {{"style":"...","material":"..."}}')
    body=json.dumps({"model":VMODEL,"prompt":prompt,"images":[b64],"stream":False,
                     "format":"json","options":{"num_predict":60,"temperature":0}}).encode()
    req=urllib.request.Request(OLLAMA+"/api/generate",body,{"Content-Type":"application/json"})
    try:
        resp=json.load(urllib.request.urlopen(req,timeout=120)).get("response","{}")
        d=json.loads(resp)
        st=d.get("style"); mt=d.get("material")
        return (st if st in STYLES else None, mt if mt in MATERIALS else None)
    except Exception as e:
        return (None, None)

def main():
    images_dir, staging, out = sys.argv[1], sys.argv[2], sys.argv[3]
    limit = None
    if "--limit" in sys.argv: limit=int(sys.argv[sys.argv.index("--limit")+1])
    if "--ollama" in sys.argv:
        global OLLAMA; OLLAMA=sys.argv[sys.argv.index("--ollama")+1]
    # index images by normalized SKU (prefer C=colorway kind)
    img_by_sku={}
    for p in glob.glob(os.path.join(images_dir,"**","*.*"),recursive=True):
        if not p.lower().endswith((".jpg",".jpeg",".png")): continue
        tok=os.path.basename(p).split("_")[0]   # LIS-42003
        sku=norm_sku(tok)
        if "_C_" in p or sku not in img_by_sku: img_by_sku[sku]=p
    rows=[json.loads(l) for l in open(staging) if l.strip()]
    n=0; enriched=0; w=open(out,"w")
    for r in rows:
        sku=norm_sku(r.get("mfr_sku"))
        p=img_by_sku.get(sku)
        if p:
            hexes=hex_palette(p)
            style,material=vision_tags(p)
            r["enrich"]={"hex":hexes,"color_family":dominant_family(hexes),
                         "style":style,"material":material,"img":os.path.basename(p),
                         "engine":"local-hybrid:pillow+qwen2.5vl:7b"}
            enriched+=1
            print(f"  {sku:12} {dominant_family(hexes):14} {str(style):20} {str(material):18} {hexes[0]['hex'] if hexes else ''}",flush=True)
        w.write(json.dumps(r,ensure_ascii=False)+"\n")
        n+=1
        if limit and enriched>=limit:
            for rest in rows[n:]: w.write(json.dumps(rest,ensure_ascii=False)+"\n")
            break
    w.close()
    print(f"\nenriched {enriched} / {n} rows  → {out}")

if __name__=="__main__": main()