Spaces:
Running on Zero
Running on Zero
File size: 4,231 Bytes
dc09dcc | 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 | #!/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()
|