`visual-search/train/eval_aug.py`

59 lines
#!/usr/bin/env python3
# A/B the augmentation re-fine-tune on the REAL task: degraded-query → clean-gallery image retrieval.
# Held-out catalog images (NOT in the training cache). Gallery = clean embeddings; queries = the SAME
# images degraded (handheld-photo sim). top-1 = does the degraded query retrieve its own clean image.
# Compares base CLIP vs dw_clip_ft.pt (current) vs dw_clip_ft_aug.pt (new) → recommends the winner.
import os, io, glob, urllib.request, urllib.parse, subprocess, random
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
import torch, torch.nn.functional as F, open_clip
from PIL import Image
import finetune_aug as A   # reuse augment()

HERE = os.path.dirname(os.path.abspath(__file__))
DB = os.environ["DW_UNIFIED_DB"]; 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

sql = ("select id,image_url from vendor_catalog where left(image_url,4)='http' "
       "and pattern_name is not null and pattern_name<>'' order by random() limit 1500")
rows = [l.split('\t') for l in subprocess.check_output(['psql',DB,'-F','\t','-tA','-c',sql]).decode().splitlines() if '\t' in l]
heldout = [r for r in rows if r[0] not in trained][:N]

random.seed(0)
clean, deglabel = [], []
for idv,url in heldout:
    try:
        data=urllib.request.urlopen(urllib.request.Request(small(url),headers={'User-Agent':'e/1'}),timeout=20).read()
        im=Image.open(io.BytesIO(data)).convert('RGB'); clean.append(im); deglabel.append(A.augment(im.copy()))
    except Exception: pass
print(f"held-out pairs: {len(clean)}", 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'); model=model.to(dev).float().eval()
CL=torch.stack([pp(i) for i in clean]).to(dev); DG=torch.stack([pp(i) for i in deglabel]).to(dev)

def top1(sd):
    if sd is None: model.load_state_dict(open_clip.create_model('ViT-B-32',pretrained='laion2b_s34b_b79k').state_dict(),strict=False)
    else: model.load_state_dict(sd,strict=False)
    with torch.no_grad():
        g=F.normalize(model.encode_image(CL),dim=-1); q=F.normalize(model.encode_image(DG),dim=-1)
        sim=q@g.t(); top=sim.argmax(dim=1); acc=(top==torch.arange(len(clean),device=dev)).float().mean().item()
        # also top-5
        t5=sim.topk(5,dim=1).indices; r5=sum(1 for i in range(len(clean)) if i in t5[i]).__truediv__(len(clean))
    return acc*100, r5*100

print(f"\n{'model':<26} deg→clean top-1   top-5")
a1,a5=top1(None); print(f"{'base laion2b':<26} {a1:5.1f}%          {a5:5.1f}%")
b1,b5=top1(torch.load(os.path.join(HERE,'ckpt','dw_clip_ft.pt'),map_location='cpu')); print(f"{'dw_clip_ft (current)':<26} {b1:5.1f}%          {b5:5.1f}%")
aug=os.path.join(HERE,'ckpt','dw_clip_ft_aug.pt')
if os.path.exists(aug):
    c1,c5=top1(torch.load(aug,map_location='cpu')); print(f"{'dw_clip_ft_aug (NEW)':<26} {c1:5.1f}%          {c5:5.1f}%")
    print(f"\n★ AUG vs current: top-1 {c1:.1f}% vs {b1:.1f}%  →  {'DEPLOY (better)' if c1>b1+1 else 'KEEP CURRENT (not better)'}")
else:
    print("\n(dw_clip_ft_aug.pt not present yet — run after training finishes)")