`enrich_sangetsu.py`

122 lines
#!/usr/bin/env python3
"""
Sangetsu enricher — pattern-level vision (style/material) + per-colorway Pillow hex.
Sangetsu rows are PATTERNS holding {colorway_skus:[...], sku_images:{sku:url}}.
  style/material : qwen2.5vl ONCE per pattern (same design across colorways), fanned out
  hex/color_family: Pillow per colorway image (each colorway = a different color)

Emits out/sangetsu-enriched.jsonl — ONE line per colorway SKU:
  {"mfr_sku": "<colorway_sku>", "source":"sangetsu", "enrich": {...}}
so the viewer's ENRICH index (keyed by normSku(mfr_sku)) picks it up per colorway card.

$0 local. Downloads images to images/sangetsu/<sku>.jpg (concurrency, skip-existing, polite).
"""
import os, json, urllib.request, concurrent.futures as cf
from enrich import hex_palette, dominant_family, vision_tags, norm_sku

IMG_DIR = "images/sangetsu"
STAGING = "staging/sangetsu-staging.jsonl"
OUT = "out/sangetsu-enriched.jsonl"
MIN_IMAGE_BYTES = 200   # reject empty/error responses (HTML 404s are ~150 B)
os.makedirs(IMG_DIR, exist_ok=True)

# Vision (style/material via qwen2.5vl) is the SLOW leg (~90s+/pattern on this box).
# color_family+hex (the primary goal) come from fast Pillow and never need it.
# Default vision OFF so the color pass completes in minutes; opt in with SANGETSU_VISION=1.
VISION = os.environ.get("SANGETSU_VISION", "0") != "0"

def fetch(sku, url):
    dst = os.path.join(IMG_DIR, norm_sku(sku) + ".jpg")
    if os.path.exists(dst) and os.path.getsize(dst) > 0: return dst
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        data = urllib.request.urlopen(req, timeout=40).read()
        if len(data) < MIN_IMAGE_BYTES: return None
        with open(dst, "wb") as f: f.write(data)
        return dst
    except Exception:
        return None

def main():
    rows = [json.loads(l) for l in open(STAGING) if l.strip()]
    # 1) gather all (sku,url) and download with modest concurrency
    jobs = []
    for r in rows:
        imgs = r.get("sku_images") or {}
        for sku, url in imgs.items():
            if url: jobs.append((sku, url))
    print(f"sangetsu: {len(rows)} patterns, {len(jobs)} colorway images to fetch  (vision={'on' if VISION else 'off'})", flush=True)
    paths = {}
    done = 0
    with cf.ThreadPoolExecutor(max_workers=8) as ex:
        futs = {ex.submit(fetch, s, u): s for s, u in jobs}
        for fut in cf.as_completed(futs):
            s = futs[fut]; p = fut.result()
            if p: paths[norm_sku(s)] = p
            done += 1
            if done % 500 == 0: print(f"  fetched {done}/{len(jobs)} ({len(paths)} ok)", flush=True)
    print(f"download complete: {len(paths)} images on disk", flush=True)

    # Resume-safe: load any prior enrichment so a re-run carries it forward
    # (skips the slow vision leg for patterns already done) and NEVER regresses.
    existing = {}
    if os.path.exists(OUT):
        for l in open(OUT):
            l = l.strip()
            if not l: continue
            try:
                o = json.loads(l)
                if o.get("mfr_sku"): existing[o["mfr_sku"]] = o
            except Exception:
                pass
    print(f"carry-forward: {len(existing)} colorways already enriched", flush=True)

    TMP = OUT + ".tmp"
    def checkpoint(lines, processed):
        # Atomic write of a SUPERSET: everything processed so far PLUS any prior
        # entry not yet reprocessed. temp+rename => a mid-run death can't truncate.
        leftover = [json.dumps(o, ensure_ascii=False) for s, o in existing.items() if s not in processed]
        with open(TMP, "w") as f:
            f.write("\n".join(lines + leftover))
            if lines or leftover: f.write("\n")
        os.replace(TMP, OUT)

    # 2) enrich: vision once per pattern (rep image), Pillow per colorway
    out_lines = []; processed = set(); n_pat = 0; n_sku = 0; n_vision = 0
    for r in rows:
        skus = r.get("colorway_skus") or []
        have = [(s, paths[norm_sku(s)]) for s in skus if norm_sku(s) in paths]
        if not have: continue
        # reuse style/material from any already-enriched colorway (skip slow vision)
        style = material = None
        for s, _ in have:
            e = (existing.get(s) or {}).get("enrich") or {}
            if e.get("style") is not None or e.get("material") is not None:
                style, material = e.get("style"), e.get("material"); break
        if VISION and style is None and material is None:
            style, material = vision_tags(have[0][1]); n_vision += 1
        n_pat += 1
        for s, p in have:
            if s in existing:
                out_lines.append(json.dumps(existing[s], ensure_ascii=False))
                processed.add(s); n_sku += 1; continue
            try:
                hexes = hex_palette(p)
            except Exception:
                hexes = []
            enrich = {"hex": hexes, "color_family": dominant_family(hexes),
                      "style": style, "material": material,
                      "img": os.path.basename(p),
                      "engine": "local-hybrid:pillow+qwen2.5vl:7b(pattern-vision)"}
            out_lines.append(json.dumps({"mfr_sku": s, "source": "sangetsu", "enrich": enrich}, ensure_ascii=False))
            processed.add(s); n_sku += 1
        if n_pat % 25 == 0:
            checkpoint(out_lines, processed)
            print(f"  enriched {n_pat} patterns / {n_sku} colorway skus  ({n_vision} vision calls; last: {r.get('pattern')} -> {style}/{material})", flush=True)
    checkpoint(out_lines, processed)
    total = len(set(processed) | set(existing))
    print(f"\nDONE: {n_sku} colorways this pass across {n_pat} patterns ({n_vision} vision calls); {total} total enriched -> {OUT}", flush=True)

if __name__ == "__main__":
    main()