Spaces:
Running on Zero
Running on Zero
| """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 | |