`visual-search/train/download.py`

65 lines
#!/usr/bin/env python3
# Resumable, threaded image downloader for the DW CLIP fine-tune.
# Reads manifest.tsv, fetches each catalog image at SMALL size (Shopify CDN width=512 so we pull
# ~20KB not ~1MB), downscales to 256px, writes cache/<id>.jpg. Skips already-cached ids so it
# resumes cleanly. Bandwidth ≈ rows × ~25KB (not × 1MB) thanks to the CDN width rewrite.
import os, io, sys, time, threading, queue, urllib.request, urllib.parse
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
CACHE = os.path.join(HERE, 'cache'); os.makedirs(CACHE, exist_ok=True)
SIZE = 256; WORKERS = 16
Image.MAX_IMAGE_PIXELS = None

def small_url(u):
    # Shopify CDN honors ?width= — ask for 512 so a 256 thumbnail is crisp but the download stays tiny.
    if 'shopify' in u or 'cdn.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

def fetch(idv, url):
    out = os.path.join(CACHE, idv + '.jpg')
    if os.path.exists(out): return 'skip'
    try:
        req = urllib.request.Request(small_url(url), headers={'User-Agent': 'dwtrain/1'})
        data = urllib.request.urlopen(req, timeout=25).read()
        im = Image.open(io.BytesIO(data)).convert('RGB')
        im.thumbnail((SIZE, SIZE))
        im.save(out, 'JPEG', quality=88)
        return 'ok'
    except Exception:
        return 'err'

def main():
    rows = []
    with open(os.path.join(HERE, 'manifest.tsv')) as f:
        for line in f:
            parts = line.rstrip('\n').split('\t')
            if len(parts) >= 2 and parts[0] and parts[1].startswith('http'):
                rows.append((parts[0], parts[1]))
    q = queue.Queue()
    for r in rows: q.put(r)
    total = len(rows)
    counts = {'ok': 0, 'skip': 0, 'err': 0}
    lock = threading.Lock(); t0 = time.time()
    def worker():
        while True:
            try: idv, url = q.get_nowait()
            except queue.Empty: return
            r = fetch(idv, url)
            with lock:
                counts[r] += 1
                done = counts['ok'] + counts['skip'] + counts['err']
                if done % 500 == 0 or done == total:
                    rate = done / max(1e-6, time.time() - t0)
                    print(f"{done}/{total}  ok={counts['ok']} skip={counts['skip']} err={counts['err']}  {rate:.0f}/s", flush=True)
            q.task_done()
    ts = [threading.Thread(target=worker, daemon=True) for _ in range(WORKERS)]
    for t in ts: t.start()
    for t in ts: t.join()
    print(f"DONE  cached={counts['ok']+counts['skip']}  err={counts['err']}", flush=True)

if __name__ == '__main__':
    main()