_prep_data.py

"""Prepare per-user data bundles for the HF Space.

For each user in USERS:
  hf_demo/data/u{N}/
    ├── vector_store/
    │   ├── manifest.json
    │   ├── records.json          (with paths rewritten to thumbs/u{N}/<pid>.jpg)
    │   └── vectors.npy
    ├── events.json
    └── images.json

  hf_demo/thumbs/u{N}/<photoid>.jpg   (resized to MAX_SIZE on long edge, JPEG q85)
"""
from __future__ import annotations
import json, shutil, os, re
from pathlib import Path

from PIL import Image

REPO    = Path('/sensei-fs-3/tenants/Sensei-AdobeResearchTeam/thaon/code')
KII_DIR = REPO / 'converted_data/database/kii/yfcc_data'
HF      = Path(__file__).resolve().parent

USERS = [3, 8]    # start with these; easy to expand later

MAX_SIZE = 384    # long-edge px; Qwen2.5-VL is happy at 384–512
JPEG_Q   = 85


def rewrite_path(orig: str, uid: int) -> str:
    """Convert an absolute disk path into the Space-relative path."""
    m = re.search(r'/(\d+\.jpg)$', orig)
    if not m:
        return orig  # leave non-matching entries (e.g. "image2.jpg") alone
    return f'thumbs/u{uid}/{m.group(1)}'


def resize_one(src: Path, dst: Path):
    if dst.exists():
        return
    dst.parent.mkdir(parents=True, exist_ok=True)
    try:
        with Image.open(src) as im:
            im = im.convert('RGB')
            im.thumbnail((MAX_SIZE, MAX_SIZE), Image.LANCZOS)
            im.save(dst, 'JPEG', quality=JPEG_Q, optimize=True)
    except Exception as e:
        print(f'  ! failed to resize {src}: {e}')


def prep_user(uid: int):
    src_dir = KII_DIR / str(uid)
    out_dir = HF / 'data' / f'u{uid}'
    thumb_dir = HF / 'thumbs' / f'u{uid}'
    print(f'\n[u{uid}] {src_dir} -> {out_dir}')

    if not src_dir.exists():
        print(f'  ! source dir missing'); return

    # Copy events.json + images.json verbatim (they're small + UI uses them).
    out_dir.mkdir(parents=True, exist_ok=True)
    for fname in ('events.json', 'images.json'):
        sp, dp = src_dir / fname, out_dir / fname
        if sp.exists():
            shutil.copy(sp, dp)

    # Vector store
    vs_src = src_dir / 'vector_store'
    vs_dst = out_dir / 'vector_store'
    vs_dst.mkdir(parents=True, exist_ok=True)
    shutil.copy(vs_src / 'manifest.json', vs_dst / 'manifest.json')
    shutil.copy(vs_src / 'vectors.npy',   vs_dst / 'vectors.npy')

    # Rewrite records.json paths to be Space-relative
    records = json.loads((vs_src / 'records.json').read_text())
    needed_pids = set()
    for r in records:
        p = r.get('payload', {})
        # images list (event records)
        if 'images' in p:
            new_imgs = []
            for orig in p['images']:
                new = rewrite_path(orig, uid)
                new_imgs.append(new)
                m = re.search(r'(\d+)\.jpg$', orig)
                if m:
                    needed_pids.add(m.group(1))
            p['images'] = new_imgs
        # single path (image record)
        if 'path' in p:
            orig = p['path']
            p['path'] = rewrite_path(orig, uid)
            m = re.search(r'(\d+)\.jpg$', orig)
            if m:
                needed_pids.add(m.group(1))
    (vs_dst / 'records.json').write_text(json.dumps(records))

    # Resize the needed photos into thumbs/u{N}/
    yfcc_imgs = REPO / 'thaodata/yfcc_v3' / str(uid) / 'images'
    n_done = 0
    n_missing = 0
    for pid in sorted(needed_pids):
        src = yfcc_imgs / f'{pid}.jpg'
        if not src.exists():
            n_missing += 1
            continue
        resize_one(src, thumb_dir / f'{pid}.jpg')
        n_done += 1
    print(f'  records: {len(records)}  thumbs: {n_done} done, {n_missing} missing')

    # Report sizes
    total_mb = sum(p.stat().st_size for p in thumb_dir.rglob('*.jpg')) / 1024 / 1024
    vs_mb    = sum(p.stat().st_size for p in vs_dst.iterdir()) / 1024 / 1024
    print(f'  thumb dir: {total_mb:.1f} MB   vector store: {vs_mb:.1f} MB')


def main():
    for uid in USERS:
        prep_user(uid)
    print('\nDONE.')


if __name__ == '__main__':
    main()