`visual-search/train/sanity_check.py`

62 lines
#!/usr/bin/env python3
# Pick the best fine-tuned checkpoint by HELD-OUT image→caption retrieval accuracy (recall@1).
# Held-out = catalog rows NOT in the training cache. Compares base CLIP vs each ckpt/dw_clip_ft_ep*.pt.
# Higher held-out recall@1 = better generalization (guards against the low-loss overfit trap).
import os, io, glob, urllib.request, urllib.parse, subprocess
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
import torch, torch.nn.functional as F, open_clip
from PIL import Image
import finetune as T   # reuse caption()

HERE = os.path.dirname(os.path.abspath(__file__))
DB = os.environ.get('DW_UNIFIED_DB', 'postgresql://dw_admin@127.0.0.1:5432/dw_unified')
N = 300
trained = set(os.path.splitext(f)[0] for f in os.listdir(os.path.join(HERE, 'cache')))

def small(u):
    if 'shopify' in u:
        pr=urllib.parse.urlparse(u); q=urllib.parse.parse_qs(pr.query); q['width']=['512']
        return urllib.parse.urlunparse(pr._replace(query=urllib.parse.urlencode(q,doseq=True)))
    return u

# pull held-out rows (id, url, pattern, color, vendor, type) not already trained
sql = ("select id,image_url,regexp_replace(pattern_name,E'[\\t\\n\\r]',' ','g'),"
       "coalesce(nullif(regexp_replace(coalesce(color_name,''),E'.*\"Name\":\\s*\"([^\"]+)\".*',E'\\1'),color_name),color_primary,''),"
       "coalesce(original_vendor_name,vendor_code,''),coalesce(product_type,'Wallcovering') "
       "from vendor_catalog where image_url like 'http%' and pattern_name is not null and pattern_name<>'' "
       "order by random() limit 1200")
rows = [l.split('\t') for l in subprocess.check_output(['psql',DB,'-F','\t','-tA','-c',sql]).decode().splitlines()]
heldout = [r for r in rows if len(r)>=6 and r[0] not in trained][:N]
print(f"held-out candidates: {len(heldout)}", flush=True)

imgs, caps = [], []
for r in heldout:
    try:
        data = urllib.request.urlopen(urllib.request.Request(small(r[1]),headers={'User-Agent':'t'}),timeout=20).read()
        imgs.append(Image.open(io.BytesIO(data)).convert('RGB')); caps.append(T.caption(r[2],r[3],r[4],r[5]))
    except Exception: pass
print(f"downloaded {len(imgs)} held-out images", flush=True)

dev='mps' if torch.backends.mps.is_available() else 'cpu'
model,_,pp = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
tok = open_clip.get_tokenizer('ViT-B-32'); model=model.to(dev).float().eval()
IM = torch.stack([pp(i) for i in imgs]).to(dev)
TX = tok(caps).to(dev)

def recall1(sd=None):
    if sd: model.load_state_dict(sd, strict=False)
    else:  model.load_state_dict(open_clip.create_model('ViT-B-32', pretrained='laion2b_s34b_b79k').state_dict(), strict=False)
    with torch.no_grad():
        imf=F.normalize(model.encode_image(IM),dim=-1); txf=F.normalize(model.encode_text(TX),dim=-1)
        sim=imf@txf.t(); top=sim.argmax(dim=1); correct=(top==torch.arange(len(imgs),device=dev)).float().mean().item()
    return correct*100

print(f"\n{'model':<26} held-out recall@1")
print(f"{'base laion2b':<26} {recall1(None):.1f}%")
best=('base',recall1(None))
for ck in sorted(glob.glob(os.path.join(HERE,'ckpt','dw_clip_ft_ep*.pt'))):
    acc=recall1(torch.load(ck,map_location='cpu')); name=os.path.basename(ck)
    print(f"{name:<26} {acc:.1f}%")
    if acc>best[1]: best=(name,acc)
print(f"\n★ BEST: {best[0]} @ {best[1]:.1f}% held-out recall@1")