Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """Match local category/scene IDs to Hub HTMLs and capture their initial cameras. | |
| Optional build tooling: npm install --no-save playwright && npx playwright install chromium. | |
| The running demo needs neither Playwright nor a browser installation. | |
| """ | |
| import argparse | |
| from concurrent.futures import ThreadPoolExecutor | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| import shutil | |
| import subprocess | |
| from urllib.parse import quote | |
| from urllib.request import urlopen | |
| ROOT = Path(__file__).resolve().parents[1] | |
| def read_json(url): | |
| with urlopen(url, timeout=60) as response: | |
| return json.load(response) | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument('--dataset', default='ritianyu/SplatAnything_Demo') | |
| parser.add_argument('--examples_dir', type=Path, default=ROOT / 'examples') | |
| parser.add_argument('--work_dir', type=Path, default=ROOT / '.cache/preview-build') | |
| args = parser.parse_args() | |
| api = f'https://huggingface.co/api/datasets/{args.dataset}' | |
| revision = read_json(api)['sha'] | |
| remote = read_json(f'{api}/tree/{revision}?recursive=true&limit=1000') | |
| scenes = {} | |
| for folder in ('object', 'indoor', 'street', 'aerial', 'rig', 'pano'): | |
| for scene in sorted((args.examples_dir / folder).iterdir()): | |
| if not scene.is_dir() or not (scene / 'images').is_dir(): | |
| continue | |
| prefix = f'{folder}/{scene.name}/' | |
| candidates = [x for x in remote if x['type'] == 'file' and x['path'].startswith(prefix) | |
| and x['path'].endswith('.html') and x['size'] > 0] | |
| if len(candidates) != 1: | |
| raise ValueError(f'Expected exactly one nonempty HTML for {prefix}: found {len(candidates)}') | |
| item = candidates[0] | |
| key = f'{folder}-{scene.name}' | |
| scenes[key] = dict(path=item['path'], bytes=item['size'], sha256=item['lfs']['oid'], poster=f'{key}.webp') | |
| catalog = dict(dataset=args.dataset, revision=revision, capture=dict(time=0, width=1440, height=900), scenes=scenes) | |
| args.work_dir.mkdir(parents=True, exist_ok=True) | |
| source = args.work_dir / 'source' | |
| source.mkdir(exist_ok=True) | |
| def fetch(item): | |
| key, record = item | |
| target = source / f'{key}.html' | |
| if target.is_file() and target.stat().st_size == record['bytes']: | |
| if hashlib.sha256(target.read_bytes()).hexdigest() == record['sha256']: | |
| return | |
| temporary = target.with_suffix('.download') | |
| try: | |
| url = f'https://huggingface.co/datasets/{args.dataset}/resolve/{revision}/' + quote(record['path'], safe='/') | |
| digest = hashlib.sha256() | |
| with urlopen(url, timeout=90) as response, temporary.open('wb') as output: | |
| while chunk := response.read(1024 * 1024): | |
| output.write(chunk) | |
| digest.update(chunk) | |
| if temporary.stat().st_size != record['bytes'] or digest.hexdigest() != record['sha256']: | |
| raise ValueError(f'Download verification failed: {key}') | |
| temporary.replace(target) | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| print(f'Downloaded {key}', flush=True) | |
| with ThreadPoolExecutor(max_workers=4) as pool: | |
| list(pool.map(fetch, scenes.items())) | |
| build_catalog = args.work_dir / 'catalog.json' | |
| build_catalog.write_text(json.dumps(catalog, indent=2) + '\n') | |
| posters = args.work_dir / 'posters' / revision | |
| subprocess.run(['node', str(ROOT / 'scripts/capture_scene_previews.cjs'), | |
| str(build_catalog.resolve()), str(source.resolve()), str(posters.resolve())], check=True) | |
| assets = ROOT / 'assets/scene-previews' | |
| assets.mkdir(parents=True, exist_ok=True) | |
| for record in scenes.values(): | |
| shutil.copyfile(posters / record['poster'], assets / record['poster']) | |
| staged = assets / 'catalog.json.tmp' | |
| staged.write_text(json.dumps(catalog, indent=2) + '\n') | |
| staged.replace(assets / 'catalog.json') | |
| print(f'Prepared {len(scenes)} initial-frame previews at {revision}. Restart the demo to use the new catalog.') | |
| if __name__ == '__main__': | |
| main() | |