Spaces:
Running on Zero
Running on Zero
File size: 3,157 Bytes
632d805 | 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 | """Per-frame vertices of a converted scene, in either layout written by data_generation/movi_to_sim_data.py:
per-file <scene>/frame_<t>/obj_<i>_pos.npy one array per object and frame (T x N files)
packed <scene>/frames.npy (T, sum_i Nv_i, 3): the same arrays, concatenated per frame
<scene>/frames_counts.npy (N,): the vertex count of every object
Both hold the same float32 values (the stored, unscaled vertices R(q) v_asset + p). The packed layout needs two files
per scene instead of ~5,000 (MOVi-B: 1,200 x 2 files instead of 3.8 million), so it fits on any local disk or a
network file system. The frame paths <scene>/frame_<t> keep naming the frames in both layouts; with the packed layout
they are names only (no folders exist).
"""
from collections import OrderedDict
from pathlib import Path
from typing import List
import numpy as np
_OPEN_LIMIT = 128 # packed scenes kept memory-mapped per process
_open = OrderedDict() # scene dir -> (frames memmap, offsets)
def is_packed(scene_dir: Path) -> bool:
return (Path(scene_dir) / 'frames.npy').is_file()
def _packed(scene_dir: Path):
key = str(scene_dir)
if key in _open:
_open.move_to_end(key)
return _open[key]
frames = np.load(Path(scene_dir) / 'frames.npy', mmap_mode='r')
counts = np.load(Path(scene_dir) / 'frames_counts.npy').astype(np.int64)
if frames.ndim != 3 or frames.shape[2] != 3 or frames.shape[1] != int(counts.sum()) or (counts <= 0).any():
raise ValueError(f'{scene_dir}: frames.npy {frames.shape} does not hold the vertex counts {counts.tolist()}')
entry = (frames, np.concatenate([[0], np.cumsum(counts)]))
_open[key] = entry
while len(_open) > _OPEN_LIMIT:
_open.popitem(last=False)
return entry
def num_frames(scene_dir: Path) -> int:
if is_packed(scene_dir):
return int(_packed(scene_dir)[0].shape[0])
return len(frame_dirs(scene_dir))
def frame_dirs(scene_dir: Path) -> List[Path]:
"""frame_<t> paths of a scene, t = 0 .. T-1 (per-file layout: the existing folders, sorted)."""
scene_dir = Path(scene_dir)
if is_packed(scene_dir):
return [scene_dir / f'frame_{t:05d}' for t in range(_packed(scene_dir)[0].shape[0])]
return sorted(p for p in scene_dir.glob('frame_*') if p.is_dir())
def load_frame(scene_dir: Path, t: int) -> List[np.ndarray]:
"""The stored vertices (Nv_i, 3) of every object at frame t, float32 (copies for the packed layout)."""
scene_dir = Path(scene_dir)
if is_packed(scene_dir):
frames, off = _packed(scene_dir)
if not 0 <= t < frames.shape[0]:
raise IndexError(f'{scene_dir}: frame {t} outside 0 .. {frames.shape[0] - 1}')
row = np.asarray(frames[t])
return [row[off[i]:off[i + 1]].copy() for i in range(len(off) - 1)]
frame_dir = scene_dir / f'frame_{t:05d}'
out = []
while (frame_dir / f'obj_{len(out)}_pos.npy').exists():
out.append(np.load(frame_dir / f'obj_{len(out)}_pos.npy'))
if not out:
raise FileNotFoundError(f'{frame_dir}: no obj_<i>_pos.npy')
return out
|