File size: 4,723 Bytes
59630ba | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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):
# Preserve the existing WorldMem 25-channel conversion exactly.
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:
# The existing dense reverse-loop evaluator uses path ordinal 7.
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
|