scripts/prepare_rgbd_train_ext_data.py
#!/usr/bin/env python3
"""Lock and materialize TRAIN-EXT, extra training scenes for the V14 data-scale test (no model).
`--stage lock` is metadata-only and freezes the ordered candidate list before any media byte is
read; `--stage prepare` downloads the locked candidates in order and applies the frozen validity
rules and the frozen camera-only frame-role rule of TRAIN-X and EVAL-V3. No model is run and no
score is read. The rule is fixed by
outputs/EXP-3D-RGBD-DATA-SCALE-V14/PLAN_BEFORE_TRAIN_EXT_DATA.md, written before this script.
TRAIN-EXT: official Hypersim train or val split (same protocol partition), not protected or
excluded, in no evaluation cohort (DEV, FRESH-V1/V2 candidates, EVAL-V3, EVAL-V4) and not in
TRAIN72, never found invalid by an earlier preparation, volume disjoint from DEV, source asset
unique against every evaluation, protected and TRAIN72 scene and within the set, cam_00 >= 16
public frames, non-BAD trajectory, in sha256(seed:scene) order; every valid scene is kept.
"""
from __future__ import annotations
import argparse
import hashlib
import importlib.util
import json
import time
import zipfile
import zlib
from pathlib import Path
import h5py
import numpy as np
import requests
from PIL import Image
from mcss.data.hypersim import convert_hypersim_pose
from mcss.mechanism_pilot.calibration import calibrated_intrinsics, published_scene_metadata
from mcss.mechanism_pilot.data_prep import BoundedRangeReader, check_geometry, resize_rgb_depth
from mcss.mechanism_pilot.optimization_bounds_contracts import FrozenTrainingPrior
from mcss.mechanism_pilot.support_redesign_data import safe_source_assets
from mcss.mechanism_pilot.unseen_data import asset_identity, csv_rows, frame_quality, select_members
EXPERIMENT = "EXP-3D-RGBD-TRAIN-EXT-DATA"
SPLIT = "TRAIN"
URL = "https://docs-assets.developer.apple.com/ml-research/datasets/hypersim/v1/scenes/{}.zip"
PLAN = Path("outputs/EXP-3D-RGBD-DATA-SCALE-V14/PLAN_BEFORE_TRAIN_EXT_DATA.md")
PLAN_SHA256 = "526c3623752317da05922d5368e341938ba9fc7b660b7e504746e910ae6eb1b6"
EVAL_V3_SCRIPT = Path("scripts/prepare_rgbd_eval_v3_data.py")
EVAL_MANIFESTS = {
"EVAL_V3": Path("outputs/EXP-3D-RGBD-EVAL-V3-DATA/manifest_eval_v3.json"),
"EVAL_V4": Path("outputs/EXP-3D-RGBD-EVAL-V4-DATA/manifest_eval_v4.json"),
}
FAILURES = {
"FRESH_SCALE_V1": Path("outputs/EXP-3D-RGBD-FRESH-SCALE-DATA-V1/scene_failures.json"),
"FRESH_V2": Path("outputs/EXP-3D-RGBD-FRESH-V2-DATA/scene_failures.json"),
"EVAL_V3": Path("outputs/EXP-3D-RGBD-EVAL-V3-DATA/scene_failures.json"),
"EVAL_V4": Path("outputs/EXP-3D-RGBD-EVAL-V4-DATA/scene_failures.json"),
}
SPLITS = ("train", "val")
TRAIN_EXT_MIN = 40
BUDGET = 2 * 1024**3
SELECTION_SEED = "TRAIN-EXT-20260930"
EXPOSURE_BASIS = (
"TRAIN-EXT is training data only: official train or val scenes outside every evaluation "
"cohort, DEV-volume disjoint and source-asset unique against every evaluation, protected and "
"TRAIN72 scene; some were observed or read before (never as an evaluation scene). The "
"protected official-test final holdout is not opened."
)
def sha(path):
return hashlib.sha256(Path(path).read_bytes()).hexdigest()
def write(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
def read(path):
return json.loads(Path(path).read_text())
def eval_v3_module(repo):
"""The frozen EVAL-V3 preparation, reused as a library (never modified)."""
spec = importlib.util.spec_from_file_location("rgbd_eval_v3_data", repo / EVAL_V3_SCRIPT)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def source_paths(repo):
paths = eval_v3_module(repo).source_paths(repo)
paths["eval_v3_preparation"] = paths.pop("preparation_source")
paths.update({f"{k.lower()}_manifest": repo / p for k, p in EVAL_MANIFESTS.items()})
paths.update({f"{k.lower()}_failures": repo / p for k, p in FAILURES.items()})
paths.update(plan=repo / PLAN, preparation_source=Path(__file__).resolve())
return paths
def selection_key(scene):
return hashlib.sha256(f"{SELECTION_SEED}:{scene}".encode()).hexdigest()
def select_candidates(
partitions, assets, allowed, trajectories, excluded_sets, blocked_assets, fence
):
"""Metadata-only candidate order; `fence` holds the forbidden (DEV) volumes."""
blocked = set(blocked_assets)
selected = []
for row in sorted(partitions, key=lambda r: selection_key(r["scene_name"])):
sid = row["scene_name"]
if (
row["official_split"] not in SPLITS
or row["protocol_partition"] != row["official_split"]
or any(sid in s for s in excluded_sets)
or sid not in assets
or sid.split("_")[1] in fence
or asset_identity(assets[sid]) in blocked
or len(allowed.get(sid, [])) < 16
or "BAD" in trajectories.get(sid + "_cam_00", {}).get("Scene type", "BAD").upper()
):
continue
selected.append(sid)
blocked.add(asset_identity(assets[sid]))
return selected
def train_ext_sets(repo, paths, partitions):
"""(excluded sets by reason, assets that block, DEV volume fence)."""
current = read(paths["current_manifest"])
v1_lock, fresh_v2_lock = read(paths["fresh_v1_lock"]), read(paths["fresh_v2_lock"])
dev = {r["scene_id"] for r in current["scenes"] if r["split"] == "DEV"}
train72 = {r["scene_id"] for r in current["scenes"] if r["split"] == "TRAIN"}
if len(dev) != 8 or len(train72) != 72:
raise ValueError("TRAIN72 + DEV manifest expected")
fresh = set(v1_lock["FRESH_V1_candidates"]) | set(fresh_v2_lock["FRESH_V2_candidates"])
evaluation = {
k: {r["scene_id"] for r in read(repo / p)["scenes"]} for k, p in EVAL_MANIFESTS.items()
}
failed = set().union(*(set(read(repo / p)) for p in FAILURES.values()))
protected = (
set(current["protected_scene_ids"])
| set(v1_lock["protected_scene_ids"])
| set(fresh_v2_lock["protected_scene_ids"])
| {
r["scene_name"]
for r in partitions
if r["protocol_partition"] in ("final_holdout", "diagnostic_test")
}
)
excluded = {next(iter(r.values())) for r in csv_rows(paths["exclusions"])}
sets = {
"protected": protected,
"excluded": excluded,
"dev": dev,
"train72": train72,
"fresh_candidates": fresh,
"eval_v3": evaluation["EVAL_V3"],
"eval_v4": evaluation["EVAL_V4"],
"previously_invalid": failed,
}
blocking = dev | train72 | fresh | evaluation["EVAL_V3"] | evaluation["EVAL_V4"] | protected
return sets, blocking, {s.split("_")[1] for s in dev}
def allowed_frames(v1, path):
allowed = {}
for split in SPLITS:
for sid, ids in v1.allowed_frames(path, split).items():
if sid in allowed:
raise ValueError("A scene appears in two official splits")
allowed[sid] = ids
return allowed
def build_lock(repo):
if sha(repo / PLAN) != PLAN_SHA256:
raise PermissionError("The V14 plan changed or is missing")
paths = source_paths(repo)
ev3 = eval_v3_module(repo)
v1 = ev3.v1_module(repo)
assets = safe_source_assets(paths["assets"])
partitions = csv_rows(paths["partitions"])
trajectories = {r["Animation"]: r for r in csv_rows(paths["trajectories"])}
sets, blocking, fence = train_ext_sets(repo, paths, partitions)
blocked_assets = {asset_identity(assets[s]) for s in blocking if s in assets}
allowed = allowed_frames(v1, paths["official_split"])
candidates = select_candidates(
partitions, assets, allowed, trajectories, sets.values(), blocked_assets, fence
)
observed = {r["scene_name"] for r in partitions if r["previously_observed"] == "True"}
return {
"schema": "mcss.rgbd_train_ext.cohort.v1",
"experiment": EXPERIMENT,
"TRAIN_EXT_candidates": candidates,
"candidate_volumes": sorted({s.split("_")[1] for s in candidates}),
"candidate_official_split": {
r["scene_name"]: r["official_split"]
for r in partitions
if r["scene_name"] in candidates
},
"previously_observed": sorted(observed & set(candidates)),
"minimum_valid": TRAIN_EXT_MIN,
"fence_volumes": sorted(fence),
"excluded_scene_ids": {k: sorted(v) for k, v in sets.items()},
"protected_scene_ids": sorted(sets["protected"]),
"assets": {s: assets[s] for s in candidates},
"rule": (
"official train or val split with the same protocol partition; not protected or "
"excluded; not DEV, TRAIN72, a FRESH-V1/V2 candidate, EVAL-V3 or EVAL-V4; not "
"invalid in an earlier preparation; volume disjoint from DEV; source asset unique vs "
"every evaluation, protected and TRAIN72 scene and within the set; cam00 >=16 public "
f"frames of its split; nonBAD trajectory; sha256('{SELECTION_SEED}:'+scene) order; "
"every candidate passing the frozen model-free validity and camera-only role rules "
"is kept (no volume cap)"
),
"exposure_basis": EXPOSURE_BASIS,
"plan_sha256": PLAN_SHA256,
"source_sha256": {str(p): sha(p) for p in paths.values()},
"data_budget_bytes": BUDGET,
"model_forward_run": False,
"protected_media_read": False,
}
def prepare(repo, root):
repo, root = Path(repo).resolve(), Path(root).resolve()
lock_path = root / "candidate_lock.json"
lock = read(lock_path)
if lock != build_lock(repo):
raise PermissionError("Candidate lock no longer reproduces from its frozen sources")
if (root / "manifest_train_ext.json").exists():
raise FileExistsError("Completed manifest_train_ext.json exists; no overwrite")
ev3 = eval_v3_module(repo)
v1 = ev3.v1_module(repo)
paths = source_paths(repo)
protected = set(lock["protected_scene_ids"])
observed = set(lock["previously_observed"])
calibration = published_scene_metadata(paths["calibration"])
bounds = read(paths["bounds"])
prior = FrozenTrainingPrior(bounds["near_m"], bounds["far_m"], bounds["source_sha256"])
camera_only_roles = v1.role_rule(repo)
allowed = allowed_frames(v1, paths["official_split"])
network = [0]
hashes, qualities, geometry, failures, supports = {}, {}, [], {}, {}
records = []
budget_stop = False
with (root / "access.jsonl").open("a") as log:
def access(sid, path, operation):
if sid in protected:
raise PermissionError("Protected scene access prohibited")
entry = {
"time": time.time(),
"scene_id": sid,
"path": str(path),
"operation": operation,
"purpose": "TRAIN_EXT_DATA_PREPARATION_NO_FORWARD",
}
log.write(json.dumps(entry) + "\n")
log.flush()
for sid in lock["TRAIN_EXT_candidates"]:
if network[0] >= BUDGET:
budget_stop = True
break
session = requests.Session()
try:
reader = BoundedRangeReader(URL.format(sid), session)
reader.shared_budget = network
with zipfile.ZipFile(reader) as archive:
try:
plan = select_members(sid, archive.infolist(), allowed[sid])
except ValueError as exc:
failures[sid] = {"status": "INELIGIBLE_MEMBERS", "error": str(exc)}
write(root / "scene_failures.json", failures)
print("INELIGIBLE", sid, "INELIGIBLE_MEMBERS", flush=True)
reader.close()
continue
plan.update(url=URL.format(sid), etag=reader.etag, split=SPLIT)
for member in plan["member_metadata"]:
name = member["name"]
dest = root / "raw" / name
access(sid, name, "download_locked_member")
if dest.exists():
content = dest.read_bytes()
if (
len(content) != member["bytes"]
or zlib.crc32(content) != member["crc32"]
):
raise ValueError("Existing download failed CRC; stop")
else:
dest.parent.mkdir(parents=True, exist_ok=True)
dest.write_bytes(archive.read(name))
hashes[str(dest)] = sha(dest)
reader.close()
finally:
session.close()
raw = root / "raw" / sid
units = {
r["parameter_name"]: r["parameter_value"]
for r in csv_rows(raw / "_detail/metadata_scene.csv")
}
scale = float(units["meters_per_asset_unit"])
if not np.isclose(
scale, float(calibration[sid]["settings_units_info_meters_scale"]), rtol=1e-6
):
failures[sid] = {"status": "INELIGIBLE_CALIBRATION_SCALE"}
write(root / "scene_failures.json", failures)
continue
def hdf(path, sid=sid):
access(sid, path, "decode_hdf5")
with h5py.File(path) as handle:
return np.array(handle["dataset"])
ids = hdf(raw / "_detail/cam_00/camera_keyframe_frame_indices.hdf5").reshape(-1)
positions = hdf(raw / "_detail/cam_00/camera_keyframe_positions.hdf5")
rotations = hdf(raw / "_detail/cam_00/camera_keyframe_orientations.hdf5")
mapping = {int(fid): j for j, fid in enumerate(ids)}
frames, quality = [], {}
prepared = root / "prepared" / sid
prepared.mkdir(parents=True, exist_ok=True)
for fid in plan["frame_ids"]:
rgb_path = raw / f"images/scene_cam_00_final_preview/frame.{fid:04d}.color.jpg"
access(sid, rgb_path, "decode_RGB")
with Image.open(rgb_path) as im:
rgb = np.asarray(im.convert("RGB"))
depth = hdf(
raw / f"images/scene_cam_00_geometry_hdf5/frame.{fid:04d}.depth_meters.hdf5"
)
if rgb.shape[:2] != depth.shape:
raise ValueError("RGB/depth dimensions differ")
quality[fid] = frame_quality(rgb, depth)
rotation = rotations[mapping[fid]]
orth = float(np.max(np.abs(rotation.T @ rotation - np.eye(3))))
det = float(np.linalg.det(rotation))
quality[fid].update(camera_orthogonality_error=orth, camera_determinant=det)
quality[fid]["eligible"] = quality[fid]["eligible"] and orth <= 0.001 and det > 0
pose = convert_hypersim_pose(
rotation, positions[mapping[fid]], meters_per_asset_unit=scale
)
if fid == plan["frame_ids"][0]:
position = hdf(
raw / f"images/scene_cam_00_geometry_hdf5/frame.{fid:04d}.position.hdf5"
)
try:
geometry.append(
check_geometry(
sid,
depth,
position,
pose,
calibrated_intrinsics(calibration[sid], depth.shape),
scale,
)
)
except ValueError as exc:
failures[sid] = {"status": "INELIGIBLE_SOURCE_GEOMETRY", "error": str(exc)}
geometry.append({"scene_id": sid, **failures[sid]})
write(root / "geometry_checks.json", geometry)
write(root / "scene_failures.json", failures)
rgb_small, depth_small = resize_rgb_depth(rgb, depth)
rgb_out, depth_out = prepared / f"{fid:04d}.png", prepared / f"{fid:04d}.npy"
Image.fromarray(rgb_small).save(rgb_out)
np.save(depth_out, depth_small)
hashes[str(rgb_out)], hashes[str(depth_out)] = sha(rgb_out), sha(depth_out)
frames.append(
{
"frame_id": fid,
"rgb": str(rgb_out),
"depth": str(depth_out),
"intrinsics": calibrated_intrinsics(calibration[sid], (128, 160)).tolist(),
"c2w": pose.tolist(),
}
)
qualities[sid] = quality
write(root / "frame_quality.json", qualities)
if sid in failures:
print("INELIGIBLE", sid, failures[sid]["status"], flush=True)
continue
try:
roles, support = camera_only_roles(sid, frames, quality, prior)
except ValueError as exc:
failures[sid] = {"status": "INELIGIBLE_FRAME_ROLES", "error": str(exc)}
write(root / "scene_failures.json", failures)
print("INELIGIBLE", sid, "INELIGIBLE_FRAME_ROLES", flush=True)
continue
records.append(
{
"scene_id": sid,
"split": SPLIT,
"historically_exposed": sid in observed,
"historical_exposure_basis": lock["exposure_basis"],
"physical_scene_id": "|".join(asset_identity(lock["assets"][sid])),
"official_split": lock["candidate_official_split"][sid],
"roles": roles,
"frames": frames,
}
)
supports[sid] = support
write(root / "role_support.json", supports)
write(root / "progress.json", {SPLIT: [r["scene_id"] for r in records]})
print("PREPARED", sid, flush=True)
write(root / "hashes.json", hashes)
write(
root / "download_cost.json",
{"network_bytes": network[0], "full_archives_downloaded": False, "budget_bytes": BUDGET},
)
status = "PASS" if len(records) >= TRAIN_EXT_MIN else "BLOCKED_INSUFFICIENT_VALID_LOCKED_SCENES"
write(
root / "manifest_train_ext.json",
{
"schema": "mcss.rgbd_train_ext.data.v1",
"image_size": [128, 160],
"depth_semantics": "ray_distance_meters",
"scenes": records,
"candidate_lock_sha256": sha(lock_path),
"protected_scene_ids": lock["protected_scene_ids"],
},
)
write(
root / "preparation_integrity.json",
{
"status": status,
"TRAIN_EXT": len(records),
"valid_volumes": sorted({r["scene_id"].split("_")[1] for r in records}),
"budget_stop": budget_stop,
"failures": failures,
"candidate_replacement": (
"none; every locked candidate is read in order until the budget"
),
"protected_media_read": False,
"model_forward_run": False,
"network_bytes": network[0],
"candidate_lock_sha256": sha(lock_path),
},
)
print(status, len(records), "failures", len(failures), flush=True)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--stage", choices=("lock", "prepare"), required=True)
parser.add_argument("--repo", type=Path, default=Path(__file__).resolve().parents[1])
parser.add_argument("--root", type=Path, default=Path("outputs") / EXPERIMENT)
args = parser.parse_args()
root = args.root.resolve()
if args.stage == "lock":
if (root / "candidate_lock.json").exists():
raise FileExistsError("Candidate lock exists; never re-lock")
lock = build_lock(args.repo.resolve())
root.mkdir(parents=True, exist_ok=True)
write(root / "candidate_lock.json", lock)
print(
json.dumps(
{
"TRAIN_EXT_candidates": len(lock["TRAIN_EXT_candidates"]),
"candidate_volumes": len(lock["candidate_volumes"]),
"previously_observed": len(lock["previously_observed"]),
}
)
)
else:
prepare(args.repo, root)
if __name__ == "__main__":
main()