Augmentation re-fine-tune: warm-start from dw_clip_ft.pt + heavy handheld-photo augmentation (rotate/crop/blur/glare/JPEG) on the image side → teach degraded→clean matching (targets the 28% real-photo top-1)

54e085f5bd38dae656f792424dd4d261889cb84a · 2026-07-07 12:57:37 -0700 · Steve Abrams
Files touched

A visual-search/train/finetune_aug.py
Diff

commit 54e085f5bd38dae656f792424dd4d261889cb84a
Author: Steve Abrams <steve@designerwallcoverings.com>
Date:   Tue Jul 7 12:57:37 2026 -0700

    Augmentation re-fine-tune: warm-start from dw_clip_ft.pt + heavy handheld-photo augmentation (rotate/crop/blur/glare/JPEG) on the image side → teach degraded→clean matching (targets the 28% real-photo top-1)
---
 visual-search/train/finetune_aug.py | 87 +++++++++++++++++++++++++++++++++++++
 1 file changed, 87 insertions(+)

diff --git a/visual-search/train/finetune_aug.py b/visual-search/train/finetune_aug.py
new file mode 100644
index 0000000..e768bdf
--- /dev/null
+++ b/visual-search/train/finetune_aug.py
@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+# AUGMENTATION re-fine-tune: warm-start from dw_clip_ft.pt and continue training with heavy
+# handheld-photo augmentation (rotate/crop/blur/glare/JPEG) on the IMAGE side, so the model learns
+# that a degraded phone shot maps to the same catalog identity — the gap robustness_test.py exposed
+# (28% top-1 on degraded photos). Caption stays clean; only the image is degraded during training.
+#   PYTORCH_ENABLE_MPS_FALLBACK=1 ./venv/bin/python finetune_aug.py --epochs 3 --batch 96
+import os, sys, io, time, random, argparse
+os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
+import torch, torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image, ImageFilter, ImageEnhance
+import open_clip
+import finetune as T   # reuse caption() + load_rows()
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+CACHE = os.path.join(HERE, 'cache'); CKPT = os.path.join(HERE, 'ckpt')
+Image.MAX_IMAGE_PIXELS = None
+
+def augment(im):
+    # simulate a handheld iPhone shot of a physical sample — varied strength, each applied ~probabilistically
+    w, h = im.size
+    if random.random() < 0.7: im = im.rotate(random.uniform(-8, 8), expand=False, fillcolor=(238,238,238))
+    if random.random() < 0.7:
+        cx, cy = random.uniform(0.02, 0.10), random.uniform(0.02, 0.10)
+        im = im.crop((int(w*cx), int(h*cy), int(w*(1-cx)), int(h*(1-cy))))
+    if random.random() < 0.6: im = im.filter(ImageFilter.GaussianBlur(random.uniform(0.4, 1.6)))
+    if random.random() < 0.7: im = ImageEnhance.Brightness(im).enhance(random.uniform(0.85, 1.30))   # glare / shadow
+    if random.random() < 0.5: im = ImageEnhance.Contrast(im).enhance(random.uniform(0.80, 1.15))
+    if random.random() < 0.6:
+        b = io.BytesIO(); im.save(b, 'JPEG', quality=random.randint(45, 80)); b.seek(0); im = Image.open(b).convert('RGB')
+    return im
+
+class AugSet(Dataset):
+    def __init__(self, rows, preprocess, tokenizer):
+        self.rows = rows; self.pp = preprocess; self.tok = tokenizer
+    def __len__(self): return len(self.rows)
+    def __getitem__(self, i):
+        idv, pat, col, ven, typ = self.rows[i]
+        im = Image.open(os.path.join(CACHE, idv + '.jpg')).convert('RGB')
+        img = self.pp(augment(im))                       # <-- degrade before CLIP preprocess
+        txt = self.tok([T.caption(pat, col, ven, typ)])[0]
+        return img, txt
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--epochs', type=int, default=3)
+    ap.add_argument('--batch', type=int, default=96)
+    ap.add_argument('--lr', type=float, default=2e-6)      # gentle — warm-starting an already-good model
+    ap.add_argument('--max_hours', type=float, default=3.0)
+    a = ap.parse_args()
+
+    dev = 'mps' if torch.backends.mps.is_available() else 'cpu'
+    model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
+    tokenizer = open_clip.get_tokenizer('ViT-B-32')
+    warm = os.path.join(CKPT, 'dw_clip_ft.pt')             # WARM-START from the current fine-tuned model
+    model.load_state_dict(torch.load(warm, map_location='cpu'), strict=False)
+    print('warm-started from dw_clip_ft.pt · device', dev, flush=True)
+    model = model.to(dev).float(); model.train()
+
+    rows = T.load_rows()
+    print(f'training pairs (cached): {len(rows)}', flush=True)
+    dl = DataLoader(AugSet(rows, preprocess, tokenizer), batch_size=a.batch, shuffle=True,
+                    num_workers=5, drop_last=True, persistent_workers=True)
+    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.1)
+    t0 = time.time(); step = 0
+    for ep in range(a.epochs):
+        for imgs, txts in dl:
+            imgs = imgs.to(dev); txts = txts.to(dev)
+            imf = F.normalize(model.encode_image(imgs), dim=-1)
+            txf = F.normalize(model.encode_text(txts), dim=-1)
+            scale = model.logit_scale.exp().clamp(max=100)
+            logits = scale * imf @ txf.t()
+            labels = torch.arange(logits.size(0), device=dev)
+            loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels)) / 2
+            opt.zero_grad(); loss.backward(); opt.step()
+            with torch.no_grad(): model.logit_scale.clamp_(0, 4.6052)
+            step += 1
+            if step % 20 == 0: print(f'ep{ep} step{step} loss={loss.item():.4f} elapsed={(time.time()-t0)/3600:.2f}h', flush=True)
+            if (time.time()-t0)/3600 > a.max_hours:
+                torch.save(model.state_dict(), os.path.join(CKPT, 'dw_clip_ft_aug.pt')); print('max_hours — saved', flush=True); return
+        torch.save(model.state_dict(), os.path.join(CKPT, f'dw_clip_ft_aug_ep{ep}.pt'))
+        print(f'saved ep{ep}', flush=True)
+    torch.save(model.state_dict(), os.path.join(CKPT, 'dw_clip_ft_aug.pt'))
+    print(f'DONE — aug weights → ckpt/dw_clip_ft_aug.pt ({(time.time()-t0)/3600:.2f}h)', flush=True)
+
+if __name__ == '__main__':
+    main()
← 6a137d9 5x contrarian gate: honest re-scope — FT visual-ID validated · back to Dw Photo Capture · eval_aug: A/B degraded→clean image retrieval (base vs curren 33641ca →