Spaces:
Running on Zero
Running on Zero
File size: 4,078 Bytes
e17a370 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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()
|