scripts/compile_dynamic_training_episodes.py
"""Compile fixed-length train/dev episodes from prepared Hypersim manifests.
The compiler reads only train and val scene manifests. It resolves and stats the RGB and
query-depth references used by each episode, but it never decodes RGB files. A small optional
train-only depth health scan is kept separate from episode compilation because it reads four
query depth arrays per eligible train scene to detect invalid supervision.
"""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
import math
from collections import Counter
from pathlib import Path, PurePosixPath
from typing import Any
import numpy as np
PROJECT = Path(__file__).resolve().parents[1]
DEFAULT_PREPARED_ROOT = PROJECT / "data/hypersim_er_prepared"
DEFAULT_PARTITIONS = PROJECT / "configs/hypersim_er_partitions.csv"
DEFAULT_OUTPUT_ROOT = PROJECT / "outputs/dynamic_training_run_20260919"
ONLINE_SCHEMA = "mcss.dynamic.online_episode.v1"
QUERY_SCHEMA = "mcss.dynamic.query_vault.v1"
INDEX_SCHEMA = "mcss.dynamic.episode_index.v1"
INVENTORY_SCHEMA = "mcss.dynamic.training_inventory.v1"
SPLIT_MAPPING = {"train": "train", "val": "dev"}
REQUIRED_COUNTS = {"train": 365, "val": 46}
EPISODE_LENGTH = 16
WARMUP_LENGTH = 4
STREAM_LENGTH = 8
QUERY_LENGTH = 4
RIGID_CAMERA_TOLERANCE = 1e-3
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _read_json(path: Path) -> dict[str, Any]:
try:
value = json.loads(path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as error:
raise ValueError(f"invalid manifest JSON: {path}") from error
if not isinstance(value, dict):
raise ValueError(f"manifest must contain an object: {path}")
return value
def _identifier(value: Any, name: str) -> str:
if not isinstance(value, str) or not value or any(char in value for char in "\\/"):
raise ValueError(f"{name} must be a simple non-empty identifier")
return value
def _image_size(value: Any, path: Path) -> list[int]:
if (
not isinstance(value, list)
or len(value) != 2
or any(type(item) is not int or item < 2 for item in value)
):
raise ValueError(f"invalid image_size in {path}")
return list(value)
def _matrix(value: Any, shape: tuple[int, int], name: str, path: Path) -> list[list[float]]:
if not isinstance(value, list) or len(value) != shape[0]:
raise ValueError(f"invalid {name} shape in {path}")
result: list[list[float]] = []
for row in value:
if not isinstance(row, list) or len(row) != shape[1]:
raise ValueError(f"invalid {name} shape in {path}")
converted = []
for item in row:
if type(item) not in {int, float} or not math.isfinite(float(item)):
raise ValueError(f"invalid {name} values in {path}")
converted.append(float(item))
result.append(converted)
if shape == (4, 4) and any(
not math.isclose(actual, expected, abs_tol=1e-5)
for actual, expected in zip(result[3], [0.0, 0.0, 0.0, 1.0], strict=True)
):
raise ValueError(f"c2w must have a homogeneous final row in {path}")
return result
def _safe_manifest_reference(scene_root: Path, value: Any, expected: Path, label: str) -> Path:
if not isinstance(value, str) or not value:
raise ValueError(f"{label} must be a non-empty relative path")
normalized = value.replace("\\", "/")
pure = PurePosixPath(normalized)
if pure.is_absolute() or any(part in {"", ".", ".."} for part in pure.parts):
raise ValueError(f"unsafe {label} path: {value!r}")
if len(pure.parts[0]) >= 2 and pure.parts[0][1] == ":":
raise ValueError(f"drive-qualified {label} path: {value!r}")
resolved = (scene_root / Path(*pure.parts)).resolve()
expected_resolved = expected.resolve()
try:
resolved.relative_to(scene_root.resolve())
except ValueError as error:
raise ValueError(f"{label} escapes scene root: {value!r}") from error
if resolved != expected_resolved:
raise ValueError(f"{label} does not match scene/frame identity: {value!r}")
try:
resolved.stat()
except OSError as error:
raise FileNotFoundError(f"missing {label} reference: {resolved}") from error
if not resolved.is_file():
raise ValueError(f"{label} reference is not a regular file: {resolved}")
return resolved
def _frame_records(manifest: dict[str, Any], path: Path) -> list[dict[str, Any]]:
raw_frames = manifest.get("frames")
if not isinstance(raw_frames, list) or not raw_frames:
raise ValueError(f"manifest frames must be a non-empty list: {path}")
frames: list[dict[str, Any]] = []
for raw in raw_frames:
if not isinstance(raw, dict):
raise ValueError(f"manifest frame must be an object: {path}")
frame_id = raw.get("frame_id")
if type(frame_id) is not int or frame_id < 0:
raise ValueError(f"invalid frame_id in {path}")
_matrix(raw.get("intrinsics"), (3, 3), "intrinsics", path)
_matrix(raw.get("c2w"), (4, 4), "c2w", path)
frames.append(raw)
frames.sort(key=lambda frame: frame["frame_id"])
ids = [frame["frame_id"] for frame in frames]
if len(ids) != len(set(ids)):
raise ValueError(f"duplicate frame_id values in {path}")
return frames
def _camera_rigid_status(
frame: dict[str, Any], *, tolerance: float = RIGID_CAMERA_TOLERANCE
) -> tuple[bool, str | None]:
rotation = np.asarray(frame["c2w"], dtype=np.float64)[:3, :3]
gram = rotation.T @ rotation
determinant = float(np.linalg.det(rotation))
if not np.allclose(gram, np.eye(3), atol=tolerance, rtol=0.0):
return False, "rotation_orthogonality"
if not math.isclose(determinant, 1.0, abs_tol=tolerance, rel_tol=0.0):
return False, "rotation_determinant"
return True, None
def _filter_camera_frames(
frames: list[dict[str, Any]], *, rigid_cameras_only: bool
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
if not rigid_cameras_only:
return frames, []
retained: list[dict[str, Any]] = []
excluded: list[dict[str, Any]] = []
for frame in frames:
valid, reason = _camera_rigid_status(frame)
if valid:
retained.append(frame)
else:
excluded.append({"frame_id": frame["frame_id"], "reason": reason})
return retained, excluded
def _load_partitions(path: Path) -> tuple[dict[str, dict[str, str]], str]:
try:
with path.open(encoding="utf-8-sig", newline="") as handle:
rows = list(csv.DictReader(handle))
except OSError as error:
raise ValueError(f"cannot read partitions CSV: {path}") from error
required = {
"scene_name",
"official_split",
"protocol_partition",
"selected_camera",
"frame_count",
}
if not rows or any(not required.issubset(row) for row in rows):
raise ValueError(f"partitions CSV is missing required columns: {path}")
selected: dict[str, dict[str, str]] = {}
for row in rows:
partition = row["protocol_partition"]
if partition not in SPLIT_MAPPING:
continue
scene = _identifier(row["scene_name"], "scene_name")
if scene in selected:
raise ValueError(f"duplicate train/val partition row: {scene}")
_identifier(row["selected_camera"], "selected_camera")
try:
int(row["frame_count"])
except ValueError as error:
raise ValueError(f"invalid partition frame_count: {scene}") from error
selected[scene] = row
return selected, sha256(path)
def _parse_scene(
manifest_path: Path,
source_split: str,
partitions: dict[str, dict[str, str]],
*,
rigid_cameras_only: bool = False,
) -> dict[str, Any]:
manifest = _read_json(manifest_path)
if manifest.get("schema_version") != 1:
raise ValueError(f"unsupported manifest schema: {manifest_path}")
if manifest.get("coordinate_convention") != "opencv_c2w":
raise ValueError(f"unsupported coordinate convention: {manifest_path}")
scene_id = _identifier(manifest.get("scene_id"), "scene_id")
if scene_id != manifest_path.parent.name:
raise ValueError(f"manifest scene_id does not match directory: {manifest_path}")
row = partitions.get(scene_id)
if row is None or row["protocol_partition"] != source_split:
raise ValueError(f"partition mismatch for {scene_id}")
if row["official_split"] != source_split:
raise ValueError(f"official split mismatch for {scene_id}")
frames = _frame_records(manifest, manifest_path)
original_frame_count = len(frames)
if int(row["frame_count"]) != original_frame_count:
raise ValueError(f"partition frame_count mismatch for {scene_id}")
frames, excluded_frames = _filter_camera_frames(
frames, rigid_cameras_only=rigid_cameras_only
)
return {
"scene_id": scene_id,
"source_split": source_split,
"split_id": SPLIT_MAPPING[source_split],
"camera": _identifier(row["selected_camera"], "selected_camera"),
"manifest_path": manifest_path.resolve(),
"manifest_sha256": sha256(manifest_path),
"image_size": _image_size(manifest.get("image_size"), manifest_path),
"frame_count": len(frames),
"original_frame_count": original_frame_count,
"rigid_cameras_only": rigid_cameras_only,
"excluded_camera_frames": excluded_frames,
"frames": frames,
}
def _window_starts(
frame_count: int,
*,
total_length: int,
stride: int,
max_windows: int | None,
) -> list[int]:
if total_length < 1 or stride < 1:
raise ValueError("window total length and stride must be positive")
if max_windows is not None and (type(max_windows) is not int or max_windows < 1):
raise ValueError("max_windows must be positive or None")
if frame_count < total_length:
return []
starts = list(range(0, frame_count - total_length + 1, stride))
return starts if max_windows is None else starts[:max_windows]
def _make_episode(
scene: dict[str, Any],
*,
start: int = 0,
warmup_steps: int = WARMUP_LENGTH,
stream_steps: int = STREAM_LENGTH,
query_count: int = QUERY_LENGTH,
window_index: int = 0,
include_window_suffix: bool = False,
) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
total_length = warmup_steps + stream_steps + query_count
if min(warmup_steps, stream_steps, query_count) < 1:
raise ValueError("warmup_steps, stream_steps, and query_count must be positive")
if start < 0 or scene["frame_count"] < start + total_length:
raise ValueError(
f"too few frames for {scene['scene_id']} window {start}:{start + total_length}"
)
selected = scene["frames"][start : start + total_length]
if len(selected) != total_length:
raise ValueError("selected window does not contain the requested number of frames")
records: list[dict[str, Any]] = []
scene_root = scene["manifest_path"].parent
image_size = scene["image_size"]
for position, frame in enumerate(selected):
frame_id = frame["frame_id"]
rgb_path = _safe_manifest_reference(
scene_root,
frame.get("rgb"),
scene_root / "rgb" / f"{frame_id:06d}.png",
f"{scene['scene_id']} frame {frame_id} rgb",
)
record = {
"frame_id": frame_id,
"intrinsics": _matrix(
frame.get("intrinsics"), (3, 3), "intrinsics", scene["manifest_path"]
),
"c2w": _matrix(frame.get("c2w"), (4, 4), "c2w", scene["manifest_path"]),
"rgb": rgb_path.as_posix(),
}
if position >= warmup_steps + stream_steps:
depth_path = _safe_manifest_reference(
scene_root,
frame.get("depth"),
scene_root / "depth" / f"{frame_id:06d}.npy",
f"{scene['scene_id']} frame {frame_id} depth",
)
record["depth"] = depth_path.as_posix()
records.append(record)
episode_id = f"{scene['split_id']}--{scene['scene_id']}--{scene['camera']}"
if include_window_suffix:
episode_id = f"{episode_id}--window-{window_index:04d}"
query_records = records[-query_count:]
query_vault_id = "qv-" + _digest({"episode_id": episode_id, "query": query_records})[:24]
online = {
"episode_id": episode_id,
"scene_id": scene["scene_id"],
"split_id": scene["split_id"],
"query_vault_id": query_vault_id,
"image_size": image_size,
"declared_length": warmup_steps + stream_steps,
"warmup": records[:warmup_steps],
"stream": records[warmup_steps : warmup_steps + stream_steps],
}
query = {
"episode_id": episode_id,
"scene_id": scene["scene_id"],
"split_id": scene["split_id"],
"query_vault_id": query_vault_id,
"image_size": image_size,
"query": query_records,
}
scene_inventory = {
"scene_id": scene["scene_id"],
"source_split": scene["source_split"],
"split_id": scene["split_id"],
"camera": scene["camera"],
"manifest": scene["manifest_path"].as_posix(),
"source_manifest_sha256": scene["manifest_sha256"],
"frame_count": scene["frame_count"],
"original_frame_count": scene["original_frame_count"],
"rigid_cameras_only": scene["rigid_cameras_only"],
"excluded_camera_frames": scene["excluded_camera_frames"],
"image_size": image_size,
"selected_frame_ids": [frame["frame_id"] for frame in selected],
"episode_id": episode_id,
"window_index": window_index,
"window_start": start,
"window_length": total_length,
"eligible": True,
"issue": None,
}
return online, query, scene_inventory
def _digest(value: Any) -> str:
canonical = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
def _scan_scenes(
prepared_root: Path,
source_split: str,
partitions: dict[str, dict[str, str]],
expected_count: int | None,
*,
rigid_cameras_only: bool = False,
) -> list[dict[str, Any]]:
split_root = (prepared_root / source_split).resolve()
if not split_root.is_dir():
raise FileNotFoundError(f"prepared split directory is missing: {split_root}")
scene_dirs = sorted(
(path for path in split_root.iterdir() if path.is_dir()), key=lambda path: path.name
)
manifest_paths: list[Path] = []
for scene_dir in scene_dirs:
manifest_path = scene_dir / "manifest.json"
if not manifest_path.is_file():
raise FileNotFoundError(f"missing scene manifest: {manifest_path}")
manifest_paths.append(manifest_path)
if expected_count is not None and len(manifest_paths) != expected_count:
raise ValueError(
f"expected {expected_count} {source_split} manifests, "
f"found {len(manifest_paths)}"
)
partition_scenes = {
scene for scene, row in partitions.items() if row["protocol_partition"] == source_split
}
scene_names = {path.parent.name for path in manifest_paths}
if scene_names != partition_scenes:
missing = sorted(partition_scenes - scene_names)
extra = sorted(scene_names - partition_scenes)
raise ValueError(
f"{source_split} manifest/partition mismatch: "
f"missing={missing}, extra={extra}"
)
return [
_parse_scene(
path,
source_split,
partitions,
rigid_cameras_only=rigid_cameras_only,
)
for path in manifest_paths
]
def _train_depth_health(
scenes: list[dict[str, Any]],
*,
warmup_steps: int,
stream_steps: int,
query_count: int,
train_window_stride: int,
max_train_windows_per_scene: int | None,
) -> dict[str, Any]:
total_length = warmup_steps + stream_steps + query_count
eligible = [scene for scene in scenes if scene["frame_count"] >= total_length]
total_values = 0
finite_positive_values = 0
shape_mismatch_frames = 0
shape_mismatch_scenes: set[str] = set()
nan_scenes: set[str] = set()
all_invalid_scenes: set[str] = set()
scanned_frames = 0
episode_count = 0
for scene in eligible:
scene_positive = 0
scene_had_nonfinite = False
starts = _window_starts(
scene["frame_count"],
total_length=total_length,
stride=train_window_stride,
max_windows=max_train_windows_per_scene,
)
episode_count += len(starts)
for start in starts:
for frame in scene["frames"][
start + warmup_steps + stream_steps : start + total_length
]:
frame_id = frame["frame_id"]
depth_path = _safe_manifest_reference(
scene["manifest_path"].parent,
frame.get("depth"),
scene["manifest_path"].parent / "depth" / f"{frame_id:06d}.npy",
f"{scene['scene_id']} frame {frame_id} depth",
)
try:
values = np.load(depth_path, allow_pickle=False)
except Exception as error:
raise ValueError(f"cannot load train query depth: {depth_path}") from error
scanned_frames += 1
array = np.asarray(values)
if array.shape != tuple(scene["image_size"]):
shape_mismatch_frames += 1
shape_mismatch_scenes.add(scene["scene_id"])
continue
finite = np.isfinite(array)
positive = finite & (array > 0)
total_values += int(array.size)
finite_positive_values += int(positive.sum())
scene_positive += int(positive.sum())
scene_had_nonfinite |= bool((~finite).any())
if scene_had_nonfinite:
nan_scenes.add(scene["scene_id"])
if scene_positive == 0:
all_invalid_scenes.add(scene["scene_id"])
return {
"scope": "train only, query depth arrays per eligible compiled window",
"scenes_available": len(scenes),
"scenes_scanned": len(eligible),
"episodes_scanned": episode_count,
"scenes_skipped_for_insufficient_frames": len(scenes) - len(eligible),
"frames_scanned": scanned_frames,
"shape_mismatch_frames": shape_mismatch_frames,
"shape_mismatch_scenes": sorted(shape_mismatch_scenes),
"nan_scenes": sorted(nan_scenes),
"all_invalid_scenes": sorted(all_invalid_scenes),
"total_values": total_values,
"finite_positive_values": finite_positive_values,
"finite_positive_ratio": (
finite_positive_values / total_values if total_values else None
),
"dev_labels_read": False,
"diagnostic_test_labels_read": False,
"final_holdout_labels_read": False,
}
def _scene_inventory(
scenes: list[dict[str, Any]], compiled: dict[str, list[dict[str, Any]]], *, total_length: int
) -> list[dict[str, Any]]:
result = []
for scene in scenes:
scene_id = scene["scene_id"]
if scene_id in compiled:
entries = compiled[scene_id]
if len(entries) == 1:
result.append(entries[0])
else:
primary = dict(entries[0])
primary["episodes"] = entries
primary["episode_count"] = len(entries)
primary["selected_frame_ids"] = []
result.append(primary)
continue
result.append(
{
"scene_id": scene_id,
"source_split": scene["source_split"],
"split_id": scene["split_id"],
"camera": scene["camera"],
"manifest": scene["manifest_path"].as_posix(),
"source_manifest_sha256": scene["manifest_sha256"],
"frame_count": scene["frame_count"],
"original_frame_count": scene["original_frame_count"],
"rigid_cameras_only": scene["rigid_cameras_only"],
"excluded_camera_frames": scene["excluded_camera_frames"],
"image_size": scene["image_size"],
"selected_frame_ids": [],
"episode_id": None,
"eligible": False,
"issue": f"insufficient_frames:{scene['frame_count']}<{total_length}",
"episode_count": 0,
}
)
return result
def _split_summary(entries: list[dict[str, Any]], expected_count: int | None) -> dict[str, Any]:
eligible_episode_count = sum(
int(entry.get("episode_count", 1)) for entry in entries if entry["eligible"]
)
return {
"manifest_count": len(entries),
"expected_manifest_count": expected_count,
"eligible_episode_count": eligible_episode_count,
"ineligible_scene_count": sum(not entry["eligible"] for entry in entries),
"frame_count_distribution": dict(
sorted(Counter(str(entry["frame_count"]) for entry in entries).items())
),
"image_size_counts": dict(
sorted(
Counter(
f"{entry['image_size'][0]}x{entry['image_size'][1]}" for entry in entries
).items()
)
),
"camera_counts": dict(sorted(Counter(entry["camera"] for entry in entries).items())),
}
def _write_json_exclusive(path: Path, value: dict[str, Any]) -> None:
with path.open("x", encoding="utf-8", newline="\n") as handle:
json.dump(value, handle, indent=2, sort_keys=True)
handle.write("\n")
def compile_training_episodes(
*,
prepared_root: str | Path = DEFAULT_PREPARED_ROOT,
partitions_path: str | Path = DEFAULT_PARTITIONS,
output_root: str | Path,
allow_ineligible: bool = False,
train_depth_health_scan: bool = True,
expected_counts: dict[str, int] | None = REQUIRED_COUNTS,
warmup_steps: int = WARMUP_LENGTH,
stream_steps: int = STREAM_LENGTH,
query_count: int = QUERY_LENGTH,
train_window_stride: int | None = None,
max_train_windows_per_scene: int | None = 1,
train_max_windows_per_scene: int | None = None,
rigid_cameras_only: bool = False,
) -> dict[str, Any]:
if min(warmup_steps, stream_steps, query_count) < 1:
raise ValueError("warmup_steps, stream_steps, and query_count must be positive")
total_length = warmup_steps + stream_steps + query_count
if train_window_stride is None:
train_window_stride = total_length
if train_window_stride < 1:
raise ValueError("train_window_stride must be positive")
if train_max_windows_per_scene is not None:
if max_train_windows_per_scene != 1:
raise ValueError(
"pass only one of max_train_windows_per_scene and train_max_windows_per_scene"
)
max_train_windows_per_scene = train_max_windows_per_scene
if max_train_windows_per_scene is not None and max_train_windows_per_scene < 1:
raise ValueError("max_train_windows_per_scene must be positive or None")
prepared = Path(prepared_root).resolve()
partitions = Path(partitions_path).resolve()
output = Path(output_root).resolve()
partition_rows, partitions_sha256 = _load_partitions(partitions)
train_scenes = _scan_scenes(
prepared,
"train",
partition_rows,
None if expected_counts is None else expected_counts["train"],
rigid_cameras_only=rigid_cameras_only,
)
dev_scenes = _scan_scenes(
prepared,
"val",
partition_rows,
None if expected_counts is None else expected_counts["val"],
rigid_cameras_only=rigid_cameras_only,
)
train_ids = {scene["scene_id"] for scene in train_scenes}
dev_ids = {scene["scene_id"] for scene in dev_scenes}
intersection = sorted(train_ids & dev_ids)
if intersection:
raise ValueError(f"train/dev scene intersection is not empty: {intersection}")
compiled_by_split: dict[str, list[tuple[dict[str, Any], dict[str, Any], dict[str, Any]]]] = {
"train": [],
"val": [],
}
compiled_inventory: dict[str, list[dict[str, Any]]] = {}
ineligible: list[str] = []
for source_split, scenes in (("train", train_scenes), ("val", dev_scenes)):
for scene in scenes:
if scene["frame_count"] < total_length:
ineligible.append(scene["scene_id"])
continue
if source_split == "train":
starts = _window_starts(
scene["frame_count"],
total_length=total_length,
stride=train_window_stride,
max_windows=max_train_windows_per_scene,
)
else:
# The development selection remains one frozen first window regardless
# of train sampling settings.
starts = [0]
include_suffix = len(starts) > 1 or starts[0] != 0
scene_entries: list[dict[str, Any]] = []
for window_index, start in enumerate(starts):
online, query, inventory = _make_episode(
scene,
start=start,
warmup_steps=warmup_steps,
stream_steps=stream_steps,
query_count=query_count,
window_index=window_index,
include_window_suffix=include_suffix,
)
compiled_by_split[source_split].append((online, query, inventory))
scene_entries.append(inventory)
compiled_inventory[scene["scene_id"]] = scene_entries
if ineligible and not allow_ineligible:
raise ValueError(
f"{total_length}-frame compilation has ineligible scenes; rerun with "
"--allow-ineligible only "
f"after accepting the explicit exclusion: {sorted(ineligible)}"
)
health = (
_train_depth_health(
train_scenes,
warmup_steps=warmup_steps,
stream_steps=stream_steps,
query_count=query_count,
train_window_stride=train_window_stride,
max_train_windows_per_scene=max_train_windows_per_scene,
)
if train_depth_health_scan
else {"enabled": False, "dev_labels_read": False, "final_holdout_labels_read": False}
)
episodes_dir = output / "episodes"
inventory_dir = output / "inventory"
planned_paths = [inventory_dir / "run_inventory.json"]
for source_split, episodes in compiled_by_split.items():
split_id = SPLIT_MAPPING[source_split]
planned_paths.append(episodes_dir / f"episode_index_{split_id}.json")
for index in range(len(episodes)):
planned_paths.extend(
(
episodes_dir / f"{split_id}_episode_{index:03d}.online.json",
episodes_dir / f"{split_id}_episode_{index:03d}.query.json",
)
)
planned_paths.append(inventory_dir / f"scene_inventory_{split_id}.json")
if any(path.exists() for path in planned_paths):
raise FileExistsError("refusing to overwrite dynamic training artifacts")
episodes_dir.mkdir(parents=True, exist_ok=True)
inventory_dir.mkdir(parents=True, exist_ok=True)
index_by_split: dict[str, list[dict[str, Any]]] = {}
for source_split, episodes in compiled_by_split.items():
split_id = SPLIT_MAPPING[source_split]
index: list[dict[str, Any]] = []
for position, (online, query, _) in enumerate(episodes):
online_path = episodes_dir / f"{split_id}_episode_{position:03d}.online.json"
query_path = episodes_dir / f"{split_id}_episode_{position:03d}.query.json"
_write_json_exclusive(online_path, {"schema_version": ONLINE_SCHEMA, **online})
_write_json_exclusive(query_path, {"schema_version": QUERY_SCHEMA, **query})
index.append(
{
"episode_id": online["episode_id"],
"scene_id": online["scene_id"],
"split_id": online["split_id"],
"query_vault_id": online["query_vault_id"],
"online_manifest": str(online_path.resolve()),
"query_manifest": str(query_path.resolve()),
"stream_steps": len(online["stream"]),
"query_count": len(query["query"]),
}
)
index_by_split[split_id] = index
_write_json_exclusive(
episodes_dir / f"episode_index_{split_id}.json",
{"schema_version": INDEX_SCHEMA, "episodes": index},
)
train_inventory = _scene_inventory(
train_scenes, compiled_inventory, total_length=total_length
)
dev_inventory = _scene_inventory(dev_scenes, compiled_inventory, total_length=total_length)
_write_json_exclusive(
inventory_dir / "scene_inventory_train.json",
{"schema_version": INVENTORY_SCHEMA, "split_id": "train", "scenes": train_inventory},
)
_write_json_exclusive(
inventory_dir / "scene_inventory_dev.json",
{"schema_version": INVENTORY_SCHEMA, "split_id": "dev", "scenes": dev_inventory},
)
inventory = {
"schema_version": INVENTORY_SCHEMA,
"status": "complete",
"prepared_root": prepared.as_posix(),
"partitions_csv": partitions.as_posix(),
"partitions_sha256": partitions_sha256,
"output_root": output.as_posix(),
"source_manifest_sha256_recorded": True,
"metadata_only_episode_preflight": True,
"rgb_files_decoded": False,
"normal_references_used": False,
"bounds_used": False,
"allow_ineligible": allow_ineligible,
"rigid_cameras_only": rigid_cameras_only,
"rigid_camera_tolerance": RIGID_CAMERA_TOLERANCE,
"rigid_camera_exclusions": [
{
"scene_id": scene["scene_id"],
"source_split": scene["source_split"],
"split_id": scene["split_id"],
"manifest": scene["manifest_path"].as_posix(),
"source_manifest_sha256": scene["manifest_sha256"],
"original_frame_count": scene["original_frame_count"],
"excluded_frame_ids": [
item["frame_id"] for item in scene["excluded_camera_frames"]
],
"excluded_frames": scene["excluded_camera_frames"],
}
for scene in [*train_scenes, *dev_scenes]
if scene["excluded_camera_frames"]
],
"window_spec": {
"warmup_steps": warmup_steps,
"stream_steps": stream_steps,
"query_count": query_count,
"total_length": total_length,
"train_window_stride": train_window_stride,
"max_train_windows_per_scene": max_train_windows_per_scene,
"dev_window_start": 0,
"dev_max_windows_per_scene": 1,
},
"ineligible_scene_ids": sorted(ineligible),
"scene_intersection": intersection,
"splits": {
"train": _split_summary(
train_inventory,
None if expected_counts is None else expected_counts["train"],
),
"dev": _split_summary(
dev_inventory,
None if expected_counts is None else expected_counts["val"],
),
},
"episode_index_train": (episodes_dir / "episode_index_train.json").resolve().as_posix(),
"episode_index_dev": (episodes_dir / "episode_index_dev.json").resolve().as_posix(),
"health_scan": health,
}
_write_json_exclusive(inventory_dir / "run_inventory.json", inventory)
return inventory
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--prepared-root", type=Path, default=DEFAULT_PREPARED_ROOT)
parser.add_argument("--partitions", type=Path, default=DEFAULT_PARTITIONS)
parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
parser.add_argument(
"--allow-ineligible",
action="store_true",
help="write eligible episodes while recording scenes that cannot satisfy 16 frames",
)
parser.add_argument(
"--skip-train-depth-health-scan",
action="store_true",
help="skip the train-only query depth health scan",
)
parser.add_argument("--warmup-steps", type=int, default=WARMUP_LENGTH)
parser.add_argument("--stream-steps", type=int, default=STREAM_LENGTH)
parser.add_argument("--query-count", type=int, default=QUERY_LENGTH)
parser.add_argument(
"--train-window-stride",
type=int,
help="chronological frame stride for train windows (default: total window length)",
)
parser.add_argument(
"--max-train-windows-per-scene",
"--train-max-windows-per-scene",
dest="max_train_windows_per_scene",
type=int,
default=1,
help="maximum train windows per scene; omit for the legacy one-window default",
)
parser.add_argument(
"--all-train-windows",
action="store_true",
help="compile every eligible train window at the configured stride",
)
parser.add_argument(
"--rigid-cameras-only",
action="store_true",
help=(
"exclude frames whose metadata c2w rotation is not rigid "
"(R.T @ R ~= I and det(R) ~= +1)"
),
)
args = parser.parse_args()
try:
max_windows = None if args.all_train_windows else args.max_train_windows_per_scene
inventory = compile_training_episodes(
prepared_root=args.prepared_root,
partitions_path=args.partitions,
output_root=args.output_root,
allow_ineligible=args.allow_ineligible,
train_depth_health_scan=not args.skip_train_depth_health_scan,
warmup_steps=args.warmup_steps,
stream_steps=args.stream_steps,
query_count=args.query_count,
train_window_stride=args.train_window_stride,
max_train_windows_per_scene=max_windows,
rigid_cameras_only=args.rigid_cameras_only,
)
except Exception as error:
print(json.dumps({"status": "failed", "error": str(error)}), flush=True)
return 1
print(
json.dumps(
{
"status": inventory["status"],
"train_episodes": inventory["splits"]["train"]["eligible_episode_count"],
"dev_episodes": inventory["splits"]["dev"]["eligible_episode_count"],
"ineligible_scene_ids": inventory["ineligible_scene_ids"],
"output_root": inventory["output_root"],
}
),
flush=True,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())