`visual-search/embed_catalog.py`

91 lines
#!/usr/bin/env python3
# Batch CLIP-embed the dw_unified catalog images for visual product search.
# Downloads each vendor_catalog.image_url, embeds with open_clip ViT-B-32 (laion2b), stores the
# L2-normalized float32 vector in image_embeddings. Resumable: rows already in image_embeddings
# (success OR permanent-fail placeholder with NULL embedding) are skipped. $0 (local CPU).
#
# Usage:  DW_UNIFIED_DB=<uri> python3 embed_catalog.py [--limit N] [--batch 32] [--workers 24]
import os, sys, io, time, struct, argparse, concurrent.futures as cf
import urllib.request
import numpy as np, psycopg2
from PIL import Image
import torch, open_clip

DB = os.environ.get("DW_UNIFIED_DB") or sys.exit("set DW_UNIFIED_DB")
ap = argparse.ArgumentParser()
ap.add_argument("--limit", type=int, default=0)     # 0 = all remaining
ap.add_argument("--batch", type=int, default=32)    # CLIP forward-pass batch
ap.add_argument("--workers", type=int, default=24)  # concurrent downloads
ap.add_argument("--chunk", type=int, default=2000)  # rows fetched per DB round
args = ap.parse_args()

print("loading CLIP ViT-B-32 (laion2b)…", flush=True)
model, _, preprocess = open_clip.create_model_and_transforms("ViT-B-32", pretrained="laion2b_s34b_b79k")

import torch as _t, os as _o
_ft = _o.path.join(_o.path.dirname(_o.path.abspath(__file__)), "dw_clip_ft.pt")
if _o.path.exists(_ft):
    model.load_state_dict(_t.load(_ft, map_location="cpu"), strict=False)
    print("loaded fine-tuned DW CLIP weights:", _ft, flush=True)

model.eval()
torch.set_num_threads(max(1, os.cpu_count() - 2))
DIM = model.visual.output_dim

def fetch(row):
    vc_id, dw_sku, mfr, vc, pat, url = row
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "dw-visual/1.0"})
        with urllib.request.urlopen(req, timeout=8) as r:
            data = r.read(6_000_000)
        img = Image.open(io.BytesIO(data)).convert("RGB")
        return (vc_id, dw_sku, mfr, vc, pat, url, preprocess(img))
    except Exception:
        return (vc_id, dw_sku, mfr, vc, pat, url, None)   # permanent-fail placeholder

def main():
    conn = psycopg2.connect(DB); conn.autocommit = False
    total_done = 0; t0 = time.time()
    while True:
        cur = conn.cursor()
        cur.execute("""select v.id, v.dw_sku, v.mfr_sku, v.vendor_code, v.pattern_name, v.image_url
                       from vendor_catalog v left join image_embeddings e on e.vc_id=v.id
                       where v.image_url is not null and v.image_url<>'' and e.vc_id is null
                       limit %s""", (args.chunk,))
        rows = cur.fetchall(); cur.close()
        if not rows: break
        # download concurrently
        got = []
        with cf.ThreadPoolExecutor(max_workers=args.workers) as ex:
            got = list(ex.map(fetch, rows))
        ok = [g for g in got if g[6] is not None]
        bad = [g for g in got if g[6] is None]
        # embed in batches
        recs = []
        for i in range(0, len(ok), args.batch):
            b = ok[i:i+args.batch]
            with torch.no_grad():
                t = torch.stack([g[6] for g in b])
                v = model.encode_image(t)
                v = (v / v.norm(dim=-1, keepdim=True)).cpu().numpy().astype("float32")
            for g, emb in zip(b, v):
                recs.append((g[0], g[1], g[2], g[3], g[4], g[5], DIM, psycopg2.Binary(emb.tobytes())))
        # write successes + permanent-fail placeholders (NULL embedding) so they aren't retried
        cur = conn.cursor()
        if recs:
            cur.executemany("""insert into image_embeddings (vc_id,dw_sku,mfr_sku,vendor_code,pattern_name,image_url,dim,embedding)
                               values (%s,%s,%s,%s,%s,%s,%s,%s) on conflict (vc_id) do nothing""", recs)
        if bad:
            cur.executemany("""insert into image_embeddings (vc_id,image_url,embedding) values (%s,%s,NULL)
                               on conflict (vc_id) do nothing""", [(g[0], g[5]) for g in bad])
        conn.commit(); cur.close()
        total_done += len(rows)
        rate = total_done / max(1e-6, time.time()-t0)
        print(f"  +{len(rows)} (ok={len(ok)} fail={len(bad)}) · {total_done} done · {rate:.1f}/s", flush=True)
        if args.limit and total_done >= args.limit: break
    conn.close()
    print(f"done: {total_done} processed in {time.time()-t0:.0f}s", flush=True)

if __name__ == "__main__":
    main()