← back to Designer Wallcoverings

onboarding/sangetsu-lilycolor/scripts/enrich-full.py

321 lines

#!/usr/bin/env python3
"""Full LOCAL swatch enrichment ($0) — scaled from enrich-pilot.py.

Reuses enrich-local-hybrid's Pillow palette + qwen2.5vl on Mac2 ollama (127.0.0.1)
and the interior-designer lexicon colorway mapper. Enumerates EVERY SKU that has an
available local swatch image (Henry Lily swatches first; Sangetsu images added only
when a local image path exists), enriches each, and APPENDS one JSONL row per SKU.

RESUMABLE: on start we read the existing output file and skip any SKU already present
(any row with a "sku" key counts as done — including error/unusable rows — so a re-run
continues rather than re-processing). ROBUST: a missing/corrupt image or a per-SKU
exception logs + skips, never crashing the whole run.

STAGES ONLY — no publish, no DB write, no Shopify. Reversible.

Output schema per SKU is IDENTICAL to the pilot: sku, hex[], bg_hex, bg_name,
dominant_hex, colorway_name, vl_color_names[], styles[], material, patterns[],
imageType, description, _provider, _cost, wall_s.

Usage: enrich-full.py [out_jsonl]
  out_jsonl defaults to staging/enrichment-full-063026A.jsonl
"""
import sys, os, json, time, subprocess, urllib.request, urllib.error, fcntl

HERE = os.path.dirname(os.path.abspath(__file__))
PROJ = os.path.dirname(HERE)

# --- image sources -----------------------------------------------------------
HENRY = os.environ.get("ENRICH_HENRY_DIR", "/Volumes/Henry/dw-lily-images")
# Sangetsu local swatch dir — only used if it actually exists on disk.
SANGETSU_DIRS = [
    os.environ.get("ENRICH_SANGETSU_DIR", "/Volumes/Henry/dw-sangetsu-images"),
    os.path.join(PROJ, "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")
PALETTE_PY = os.path.expanduser("~/Projects/enrich-local-hybrid/enrich-palette.py")
LEXICON_PY = os.path.join(HERE, "lexicon-colorway.py")
DEFAULT_OUT = os.path.join(PROJ, "staging", "enrichment-full-063026A.jsonl")
IMG_EXTS = (".jpg", ".jpeg", ".png")
# How many times to wait-and-retry a single SKU through a transient ollama outage
# before deferring it (writing nothing, so a later run reclaims it). Backoff is
# capped at 60s/attempt, so 30 ≈ up to ~25 min of tolerated downtime per SKU.
MAX_TRANSIENT_RETRIES = int(os.environ.get("ENRICH_MAX_RETRIES", "30"))


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


class VolumeUnavailable(Exception):
    """The image's source volume (e.g. Henry) is not mounted — an INFRASTRUCTURE
    outage, not a missing file. Routed through the transient retry/defer path so a
    mid-run drive drop never burns a SKU as a permanent image_not_found stub."""


def enrich_one(sku, img):
    if not os.path.exists(img):
        # If the containing directory is gone, the whole volume dropped (Henry
        # unmounted) → transient, defer. If the dir is present but the file is
        # absent, it's a genuinely missing image → permanent stub.
        if not os.path.isdir(os.path.dirname(img)):
            raise VolumeUnavailable(f"volume for {img} unavailable")
        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 discover_skus():
    """Return an ordered list of (sku, image_path). Henry Lily swatches first,
    then any Sangetsu local swatches whose directory actually exists. Dedupes by
    SKU (first source wins) and sorts within each source for stable, resumable order."""
    seen = set()
    items = []

    def add_dir(d):
        if not os.path.isdir(d):
            return
        for fn in sorted(os.listdir(d)):
            stem, ext = os.path.splitext(fn)
            if ext.lower() not in IMG_EXTS or not stem:
                continue
            if stem in seen:
                continue
            seen.add(stem)
            items.append((stem, os.path.join(d, fn)))

    add_dir(HENRY)                    # Lily swatches (primary source)
    for d in SANGETSU_DIRS:           # Sangetsu only if a local dir exists
        add_dir(d)
    return items


def load_done(out_path):
    """Read the existing output and return the set of SKUs already processed."""
    done = set()
    if not os.path.exists(out_path):
        return done
    with open(out_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except Exception:
                continue  # tolerate a truncated/corrupt trailing line
            s = obj.get("sku")
            if s:
                done.add(s)
    return done


def fmt_eta(seconds):
    seconds = int(max(0, seconds))
    h, rem = divmod(seconds, 3600)
    m, s = divmod(rem, 60)
    if h:
        return f"{h}h{m:02d}m"
    if m:
        return f"{m}m{s:02d}s"
    return f"{s}s"


def main():
    out_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_OUT
    os.makedirs(os.path.dirname(out_path), exist_ok=True)

    # --- single-instance guard (per output file) -------------------------------
    # A second worker targeting the SAME shard file refuses to start, instead of
    # stacking. On 2026-07-01, re-invoking the launcher stacked 5 orphaned
    # enrich-full.py procs on shard0/shard1 that thrashed Mac2 swap to a
    # near-freeze. Distinct shard files (shard0/shard1/sweep) still run in
    # parallel because each gets its own lock. Lock releases on process exit.
    _lock_fh = open(out_path + ".lock", "w")
    try:
        fcntl.flock(_lock_fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        # Exit NON-zero (busy) so a supervisor treats this as "retry later", not
        # completion. A clean exit(0) here would make supervise.sh's "rc==0 + no
        # new rows" check false-conclude the job is done while a sibling run holds
        # the lock. 75 = EX_TEMPFAIL (try again).
        print(f"[guard] another enrich-full.py is already running for "
              f"{os.path.basename(out_path)} — exiting busy (retry later).", flush=True)
        sys.exit(75)
    _lock_fh.write(str(os.getpid()) + "\n"); _lock_fh.flush()
    globals()["_ENRICH_LOCK_FH"] = _lock_fh  # keep ref alive so lock persists

    all_items = discover_skus()
    total = len(all_items)
    done = load_done(out_path)
    # Extra "already done" files (e.g. a seed from an earlier single-stream run)
    # so shards don't redo SKUs already enriched. Colon-separated paths.
    for seed in filter(None, os.environ.get("ENRICH_DONE_SEEDS", "").split(":")):
        done |= load_done(seed)
    todo = [(sku, img) for sku, img in all_items if sku not in done]

    # Optional cross-machine sharding: this process handles only every
    # SHARD_COUNT-th remaining item (offset SHARD_INDEX). Endpoint is chosen via
    # ENRICH_OLLAMA_URL, so shard 0 -> Mac2, shard 1 -> Mac1, etc. Each shard
    # writes its own out file; merge them when done. Fully reversible.
    shard_count = int(os.environ.get("ENRICH_SHARD_COUNT", "1"))
    shard_index = int(os.environ.get("ENRICH_SHARD_INDEX", "0"))
    if shard_count > 1:
        todo = [it for n, it in enumerate(todo) if n % shard_count == shard_index]

    print(f"discovered={total} already_done={len(done)} todo={len(todo)} "
          f"out={out_path}", flush=True)
    if not todo:
        print("Nothing to do — all discovered SKUs already enriched.", flush=True)
        return

    ok = fail = deferred = 0
    t_start = time.time()
    # Append mode — never clobber prior work (resumable).
    with open(out_path, "a") as fout:
        for i, (sku, img) in enumerate(todo, 1):
            attempt = 0
            while True:
                try:
                    row = enrich_one(sku, img)
                    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()
                    break
                except (urllib.error.URLError, TimeoutError, ConnectionError,
                        VolumeUnavailable) as e:
                    # TRANSIENT outage — ollama down (Errno 61 / timeout) OR the
                    # image volume (Henry) unmounted mid-run.
                    # NEVER write a stub: a stub row carries a "sku" key, which load_done
                    # would treat as done and skip forever. Wait for ollama to return and
                    # retry the SAME SKU with capped backoff; if it stays down past the
                    # retry budget, defer (write nothing) so a later run reclaims it.
                    attempt += 1
                    if attempt > MAX_TRANSIENT_RETRIES:
                        deferred += 1
                        print(f"[{i}/{len(todo)}] {sku} DEFER after {attempt-1} retries "
                              f"(ollama still down) — left for a later sweep", flush=True)
                        break
                    wait = min(60, 5 * attempt)
                    print(f"[{i}/{len(todo)}] {sku} transient outage ({str(e)[:60]}) — "
                          f"retry {attempt}/{MAX_TRANSIENT_RETRIES} in {wait}s", flush=True)
                    time.sleep(wait)
                    continue
                except Exception as e:
                    # PERMANENT per-SKU error (bad/missing/corrupt image). Stub it so it
                    # counts as done and the run doesn't retry a genuinely broken image.
                    fail += 1
                    fout.write(json.dumps({"sku": sku, "error": str(e)[:200]}) + "\n")
                    fout.flush()
                    print(f"[{i}/{len(todo)}] {sku} ERROR {str(e)[:120]}", flush=True)
                    break

            if i % 25 == 0 or i == len(todo):
                elapsed = time.time() - t_start
                rate = i / elapsed if elapsed > 0 else 0            # skus/sec
                remaining = len(todo) - i
                eta = remaining / rate if rate > 0 else 0
                total_done = len(done) + i
                print(f"[{i}/{len(todo)}] progress {total_done}/{total} "
                      f"ok={ok} fail={fail} deferred={deferred} rate={rate*60:.1f}/min "
                      f"eta={fmt_eta(eta)}", flush=True)

    print(f"\nDONE this pass: ok={ok} fail={fail} deferred={deferred} "
          f"processed={len(todo)} grand_total_now={len(done)+len(todo)-deferred}/{total}"
          f"{' (deferred SKUs reclaimed on next run)' if deferred else ''}", flush=True)


if __name__ == "__main__":
    main()