| """Fixed WorldMem benchmark identities, RGB frames, and model-independent I/O.""" |
|
|
| import csv |
| import json |
| import re |
| from pathlib import Path |
|
|
| import av |
| import numpy as np |
| import torch |
|
|
|
|
| def select_cases(manifest, scope=None, rank=0, world_size=1, limit=None): |
| with open(manifest, newline="") as stream: |
| rows = list(csv.DictReader(stream)) |
| if scope is not None and "scope" in rows[0]: |
| rows = [row for row in rows if row["scope"] == scope] |
| if limit is not None: |
| rows = rows[:limit] |
| if not rows or not 0 <= rank < world_size: |
| raise ValueError("Empty benchmark selection or invalid worker shard") |
| return rows[rank::world_size] |
|
|
|
|
| def read_frames(path, indices): |
| indices = torch.as_tensor(indices, dtype=torch.long) |
| first, last = int(indices.min()), int(indices.max()) |
| frames = [] |
| with av.open(str(path)) as container: |
| for ordinal, frame in enumerate(container.decode(video=0)): |
| if ordinal > last: |
| break |
| if ordinal >= first: |
| frames.append(torch.from_numpy(frame.to_ndarray(format="rgb24"))) |
| if len(frames) != last - first + 1: |
| raise ValueError(f"Missing requested frames in {path}") |
| video = torch.stack(frames).index_select(0, indices - first) |
| return video.permute(0, 3, 1, 2).float().div_(255.0).contiguous() |
|
|
|
|
| def _scene_file(root, tree, scene_id, suffix): |
| directory = Path(root) / tree |
| direct = directory / f"{scene_id}{suffix}" |
| if direct.is_file(): |
| return direct |
| matches = list(directory.glob(f"*/{scene_id}{suffix}")) |
| if len(matches) != 1: |
| raise ValueError(f"Expected exactly one {tree}/{scene_id}{suffix}") |
| return matches[0] |
|
|
|
|
| def load_re10k_case(row, root): |
| start = int(row.get("history_start", 0)) |
| indices = torch.cat((torch.arange(200), torch.arange(199, -1, -1))) + start |
| video_path = _scene_file(root, "test_256", row["scene_id"], ".mp4") |
| pose_path = _scene_file(root, "test_poses", row["scene_id"], ".pt") |
| poses = torch.load(pose_path, map_location="cpu", weights_only=True)[indices].float() |
| return { |
| "rgb": read_frames(video_path, indices), |
| "raw_poses": poses, |
| "cameras": torch.cat((poses[:, :4], poses[:, 6:]), dim=-1), |
| "history_frames": 100, |
| "history": 100, |
| "source_indices": indices, |
| "video_path": str(video_path), |
| "row": dict(row), |
| } |
|
|
|
|
| def minecraft_actions(actions): |
| |
| actions = torch.as_tensor(actions) |
| result = torch.zeros((len(actions), 25), dtype=torch.float32) |
| result[actions[:, 0] == 1, 11] = 1 |
| result[actions[:, 0] == 2, 12] = 1 |
| result[actions[:, 4] == 11, 16] = -1 |
| result[actions[:, 4] == 13, 16] = 1 |
| result[actions[:, 3] == 11, 15] = -1 |
| result[actions[:, 3] == 13, 15] = 1 |
| result[(actions[:, 5] == 6) | (actions[:, 5] == 1), 24] = 1 |
| result[actions[:, 1] == 1, 13] = 1 |
| result[actions[:, 1] == 2, 14] = 1 |
| result[actions[:, 7] == 1, 2] = 1 |
| return result |
|
|
|
|
| def load_minecraft_case(row, root): |
| indices = torch.arange(int(row["context_start"]), int(row["target_end"])) |
| path = Path(root) / row["relative_path"] |
| with np.load(path.with_suffix(".npz")) as annotation: |
| actions = minecraft_actions(annotation["actions"])[indices] |
| poses = torch.from_numpy(annotation["poses"].copy()).float()[indices] |
| return { |
| "rgb": read_frames(path, indices), |
| "actions": actions, |
| "poses": poses, |
| "history_frames": int(row["context_frames"]), |
| "history": int(row["context_frames"]), |
| "source_indices": indices, |
| "video_path": str(path), |
| "row": dict(row), |
| } |
|
|
|
|
| def case_key(row): |
| if "scene_id" in row: |
| return row["scene_id"] |
| clean = re.sub(r"[^A-Za-z0-9._-]+", "__", row["clip_id"]).strip("_") |
| return f"g{int(row['global_test_index']):06d}_{clean}" |
|
|
|
|
| def case_seed(row, seed): |
| if "scene_id" in row: |
| |
| return int(seed) * 1_000_000 + 1_000_003 * int(row["metadata_index"]) + 7 |
| return int(seed) * 1_000_000 + int(row["global_test_index"]) |
|
|
|
|
| def save_case(output, row, prediction, metadata): |
| """Stage exact float predictions; the scorer removes them after video/metric export.""" |
| directory = Path(output) / "cases" / case_key(row) |
| directory.mkdir(parents=True, exist_ok=True) |
| torch.save(prediction.detach().float().cpu().contiguous(), directory / "prediction.pt") |
| (directory / "case.json").write_text( |
| json.dumps({"case": dict(row), **metadata}, indent=2, default=str) + "\n" |
| ) |
| return directory |
|
|
