src/mcss/mechanism_pilot/support_redesign_data.py
"""New support-redesign data cohorts, with sealed final holdout metadata.
Only source metadata, camera geometry, RGB/depth validity are read. No model output.
Old preparation modules and old experiment hashes are never modified.
"""
from __future__ import annotations
import ast
import json
import time
import zipfile
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.dynamic.checkpoint import load_dynamic_checkpoint
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.small_training import sha, write_json
from mcss.mechanism_pilot.unseen_data import (
asset_identity,
csv_rows,
frame_quality,
select_members,
select_roles,
validate_unseen,
)
def safe_source_assets(path):
"""Parse literal dictionaries and os.path.join of literal strings, never execute."""
rows = {}
for node in ast.walk(ast.parse(Path(path).read_text())):
if not (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and isinstance(node.func.value, ast.Name)
and node.func.value.id == "scenes"
and node.func.attr == "append"
and len(node.args) == 1
and isinstance(node.args[0], ast.Dict)
):
continue
row = {}
for key, value in zip(node.args[0].keys, node.args[0].values, strict=True):
key = ast.literal_eval(key)
if isinstance(value, ast.Call):
if ast.unparse(value.func) != "os.path.join" or value.keywords:
raise ValueError("Unsupported source metadata expression")
parts = [ast.literal_eval(a) for a in value.args]
if not all(isinstance(a, str) and not a.startswith(("/", "\\")) for a in parts):
raise ValueError("Invalid literal source asset path")
row[key] = "/".join(parts)
else:
row[key] = ast.literal_eval(value)
rows[row["name"]] = row
return rows
def validate_roles_disjoint(train, exposed, dev, holdout, assets):
groups = [set(x) for x in (train, exposed, dev, holdout)]
if any(groups[i] & groups[j] for i in range(4) for j in range(i + 1, 4)):
raise ValueError("Train/exposed/dev/final holdout overlap")
all_ids = train + exposed + dev + holdout
identities = [asset_identity(assets[s]) for s in all_ids]
if len(identities) != len(set(identities)):
raise ValueError("Source asset overlap across cohorts")
volumes = [s.split("_")[1] for s in all_ids]
if len(volumes) != len(set(volumes)):
raise ValueError("Source volume overlap across cohorts")
def prepare_support_redesign(
output, experiment, checkpoint_path, training_dir, training_manifest, partitions, calibration
):
output, experiment, training_dir = Path(output), Path(experiment), Path(training_dir)
if output.exists():
raise FileExistsError("Refusing overwrite of locked unseen cohort")
metadata = experiment / "rgb_audit/official_metadata"
upstream = experiment / "metadata/upstream_dataset_config.py"
assets = safe_source_assets(upstream)
train_manifest = json.loads(Path(training_manifest).read_text())
train_ids = sorted(s["scene_id"] for s in train_manifest["scenes"])
log_ids = sorted(
{
json.loads(s)["scene_id"]
for s in (training_dir / "training.jsonl").read_text().splitlines()
}
)
training_summary = json.loads((training_dir / "summary.json").read_text())
if (
train_ids != ["ai_001_001", "ai_002_001", "ai_003_001"]
or train_ids != log_ids
or training_summary["n_train_scenes"] != 3
):
raise ValueError("Training identities not consistent")
carrier, _, checkpoint = load_dynamic_checkpoint(checkpoint_path)
if checkpoint["sha256"] != training_summary["checkpoint_sha256"]["phase_b_final"]:
raise ValueError("Checkpoint hash mismatch")
partition_rows = csv_rows(partitions)
trajectories = {
r["Animation"]: r for r in csv_rows(metadata / "metadata_camera_trajectories.csv")
}
split_rows = csv_rows(metadata / "metadata_images_split_scene_v1.csv")
allowed = {}
for row in split_rows:
if (
row["camera_name"] == "cam_00"
and row["included_in_public_release"] == "True"
and row["split_partition_name"] == "train"
):
allowed.setdefault(row["scene_name"], []).append(int(row["frame_id"]))
old_manifest = json.loads((experiment / "unseen_data/manifest.json").read_text())
exposed_ids = sorted(s["scene_id"] for s in old_manifest["scenes"])
# Entire old candidate cohort, including geometry-ineligible ai005003, is excluded.
old_candidates = json.loads((experiment / "unseen_data/candidate_lock.json").read_text())[
"candidate_scene_ids"
]
selected, volumes = [], {s.split("_")[1] for s in train_ids + old_candidates}
for row in sorted(partition_rows, key=lambda r: r["scene_name"]):
scene = row["scene_name"]
volume = scene.split("_")[1]
if (
volume in volumes
or row["protocol_partition"] != "train"
or scene not in assets
or len(allowed.get(scene, [])) < 16
):
continue
trajectory = trajectories.get(scene + "_cam_00")
if trajectory is None or "BAD" in trajectory["Scene type"].upper():
continue
selected.append(scene)
volumes.add(volume)
if len(selected) == 24:
break
dev_ids, holdout_ids = selected[::2], selected[1::2]
validate_roles_disjoint(train_ids, exposed_ids, dev_ids, holdout_ids, assets)
prereg = output.parent / "preregistration.json"
if not prereg.exists():
raise ValueError("Unified preregistration must exist before any cohort lock/download")
validate_unseen(train_ids, selected, assets)
output.mkdir(parents=True)
source_paths = [
upstream,
Path(partitions),
metadata / "metadata_camera_trajectories.csv",
metadata / "metadata_images_split_scene_v1.csv",
Path(training_manifest),
training_dir / "summary.json",
training_dir / "training.jsonl",
Path(calibration),
]
lock = {
"schema": "mcss.support_redesign.candidates.v1",
"preregistration_sha256": sha(prereg),
"REDESIGN_DEV_SCENES": dev_ids,
"FINAL_HOLDOUT_SCENES": holdout_ids,
"REDESIGN_DEV_EXPOSED": exposed_ids,
"native_geometry_p95_limit_m": 0.01,
"holdout_model_access": "PROHIBITED_UNTIL_METHOD_LOCK_AND_DEV_GATE",
"created_unix": time.time(),
"checkpoint_sha256": checkpoint["sha256"],
"train_scene_ids": train_ids,
"candidate_scene_ids": selected,
"assets": {s: assets[s] for s in train_ids + selected},
"selection": (
"lexicographic first eligible protocol-train cam00 non-BAD scene "
"from each distinct untrained asset volume; 8 fixed candidates, no "
"replacement"
),
"physical_identity_limit": (
"Distinct official volume/archive/asset identities; vendor cross-"
"asset content reuse cannot be independently ruled out."
),
"previously_observed": {
r["scene_name"]: r["previously_observed"]
for r in partition_rows
if r["scene_name"] in selected
},
"source_sha256": {str(p): sha(p) for p in source_paths},
"upstream_url": "https://github.com/apple/ml-hypersim/blob/main/evermotion_dataset/_dataset_config.py",
"data_budget_bytes": 2 * 1024**3,
"quality_rule": (
"finite RGB, >99.9 percent pixels all channels <=2/255 rejected; "
"finite positive native depth fraction >=0.95; no scores"
),
"role_rule": (
"first quality frame anchor, last two quality frames queries; "
"enumerate disjoint pairs from remaining frames. Maximize "
"min(A_support,B_support), minimize support gap, then mean baseline "
"gap, then lex frame IDs. At least 7 quality frames; each context "
">=2 supported candidates. This minimum is not a quality proof."
),
"model_predictions_read": False,
}
write_json(output / "candidate_lock.json", lock)
write_json(
output.parent / "scene_role_lock.json",
{
"created_unix": time.time(),
"stage": "INITIAL_IMMUTABLE_COHORT_SELECTION",
"candidate_lock_sha256": sha(output / "candidate_lock.json"),
"TRAIN_SCENES": train_ids,
"REDESIGN_DEV_EXPOSED": exposed_ids,
"REDESIGN_DEV_SCENES": dev_ids,
"FINAL_HOLDOUT_SCENES": holdout_ids,
"FINAL_HOLDOUT_PROTECTED": True,
"ineligible_holdout_transfer_to_dev": "PROHIBITED",
"qualified_frame_roles_will_be_separate": "qualified_frame_role_lock.json",
"source_assets": {s: assets[s] for s in train_ids + exposed_ids + selected},
"model_predictions_read": False,
},
)
print("CANDIDATE_LOCKED " + json.dumps(selected), flush=True)
sessions, readers, archives, plans = [], [], [], []
shared_budget = [0]
calibration_rows = published_scene_metadata(calibration)
try:
for scene in selected:
session = requests.Session()
url = f"https://docs-assets.developer.apple.com/ml-research/datasets/hypersim/v1/scenes/{scene}.zip"
reader = BoundedRangeReader(url, session)
reader.shared_budget = shared_budget
archive = zipfile.ZipFile(reader)
plan = select_members(scene, archive.infolist(), allowed[scene])
plan.update(url=url, etag=reader.etag, archive_bytes=reader.size)
sessions.append(session)
readers.append(reader)
archives.append(archive)
plans.append(plan)
print("MEMBER_PLAN " + scene, flush=True)
if sum(p["compressed_bytes"] + p["uncompressed_bytes"] for p in plans) > 2 * 1024**3:
raise ValueError("Selected data exceeds 2GiB bound")
write_json(
output / "data_lock.json",
{
"created_unix": time.time(),
"candidate_lock_sha256": sha(output / "candidate_lock.json"),
"scenes": plans,
},
)
print("ALL_MEMBERS_LOCKED", flush=True)
hashes, quality_rows, feasibility, scenes, geometry = {}, {}, {}, [], []
with (output / "access.jsonl").open("x") as log:
def access(sid, path, operation):
log.write(
json.dumps(
{
"time": time.time(),
"scene_id": sid,
"path": str(path),
"operation": operation,
"purpose": "unseen_data_geometry_only",
"official_split": "train",
}
)
+ "\n"
)
log.flush()
for plan, archive, _reader in zip(plans, archives, readers, strict=True):
sid = plan["scene_id"]
for member in plan["members"]:
dest = output / "raw" / member
dest.parent.mkdir(parents=True, exist_ok=True)
access(sid, member, "download")
dest.write_bytes(archive.read(member))
hashes[str(dest.relative_to(output))] = sha(dest)
raw = output / "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_rows[sid]["settings_units_info_meters_scale"]),
rtol=1e-6,
):
raise ValueError("Native calibration scale mismatch")
def hdf(path, sid=sid):
access(sid, path, "decode_hdf5")
with h5py.File(path) as f:
return np.array(f["dataset"])
ids = hdf(raw / "_detail/cam_00/camera_keyframe_frame_indices.hdf5").reshape(-1)
poses = hdf(raw / "_detail/cam_00/camera_keyframe_positions.hdf5")
rotations = hdf(raw / "_detail/cam_00/camera_keyframe_orientations.hdf5")
mapping = {int(i): j for j, i in enumerate(ids)}
frames, quality = [], {}
prepared = output / "prepared" / sid
prepared.mkdir(parents=True)
for fid in plan["frame_ids"]:
rgbpath = raw / f"images/scene_cam_00_final_preview/frame.{fid:04d}.color.jpg"
access(sid, rgbpath, "decode_RGB")
with Image.open(rgbpath) as im:
rgb = np.asarray(im.convert("RGB"))
depth = hdf(
raw / f"images/scene_cam_00_geometry_hdf5/frame.{fid:04d}.depth_meters.hdf5"
)
quality[fid] = frame_quality(rgb, depth)
if rgb.shape[:2] != depth.shape:
raise ValueError("Native RGB/depth dimensions disagree")
pose = convert_hypersim_pose(
rotations[mapping[fid]], poses[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_rows[sid], depth.shape),
scale,
)
)
except ValueError as exc:
geometry.append({"scene_id": sid, "status": "FAIL", "error": str(exc)})
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)
for path in (rgb_out, depth_out):
hashes[str(path.relative_to(output))] = sha(path)
frames.append(
{
"frame_id": fid,
"rgb": str(rgb_out.resolve()),
"depth": str(depth_out.resolve()),
"intrinsics": calibrated_intrinsics(
calibration_rows[sid], (128, 160)
).tolist(),
"c2w": pose.tolist(),
}
)
roles, diagnostic = select_roles(
frames, quality, carrier._candidate_points.detach()
)
if any(r.get("status") == "FAIL" and r["scene_id"] == sid for r in geometry):
diagnostic.update(status="INELIGIBLE", reason="native_geometry_check_failed")
quality_rows[sid], feasibility[sid] = quality, diagnostic
if diagnostic["status"] == "ELIGIBLE":
scenes.append(
{
"scene_id": sid,
"split": "dev" if sid in dev_ids else "final_holdout",
"cohort_role": "REDESIGN_DEVELOPMENT"
if sid in dev_ids
else "FINAL_QUALIFICATION_HOLDOUT",
"physical_scene_id": "|".join(asset_identity(assets[sid])),
"frames": frames,
"roles": roles,
}
)
write_json(output / "progress_quality.json", quality_rows)
write_json(output / "progress_feasibility.json", feasibility)
print("PREPARED " + sid + " " + json.dumps(diagnostic), flush=True)
dev = [s for s in scenes if s["scene_id"] in dev_ids]
holdout = [s for s in scenes if s["scene_id"] in holdout_ids]
status = "READY" if len(dev) >= 8 and len(holdout) >= 8 else "BLOCKED_INSUFFICIENT_DATA"
common = {
"schema": "mcss.support_redesign.data.v1",
"image_size": [128, 160],
"depth_semantics": "ray_distance_meters",
"DATA_COHORT_STATUS": status,
"candidate_lock_sha256": sha(output / "candidate_lock.json"),
"data_lock_sha256": sha(output / "data_lock.json"),
"model_predictions_read": False,
}
manifest = dict(common, scenes=scenes)
write_json(output / "manifest.json", manifest)
write_json(output / "dev_manifest.json", dict(common, scenes=dev))
write_json(
output / "holdout_manifest.json",
dict(
common, scenes=holdout, inference_access="PROHIBITED_UNTIL_METHOD_LOCK_AND_DEV_GATE"
),
)
write_json(output / "frame_quality.json", quality_rows)
write_json(output / "support_feasibility.json", feasibility)
write_json(output / "geometry_checks.json", geometry)
write_json(output / "hashes.json", hashes)
role_lock = {
"created_unix": time.time(),
"manifest_sha256": sha(output / "manifest.json"),
"candidate_lock_sha256": sha(output / "candidate_lock.json"),
"TRAIN_SCENES": train_ids,
"REDESIGN_DEV_EXPOSED": exposed_ids,
"REDESIGN_DEV_SCENES": dev_ids,
"FINAL_HOLDOUT_SCENES": holdout_ids,
"eligible_redesign_dev": [s["scene_id"] for s in dev],
"eligible_final_holdout": [s["scene_id"] for s in holdout],
"roles": {s["scene_id"]: s["roles"] for s in scenes},
"DATA_COHORT_STATUS": status,
"model_predictions_read": False,
"FINAL_HOLDOUT_PROTECTED": True,
"holdout_inference_requires": "qualified dev gate AND frozen method lock",
"source_assets": {s: assets[s] for s in train_ids + exposed_ids + selected},
"dataset_hashes_sha256": sha(output / "hashes.json"),
"initial_scene_role_lock_sha256": sha(output.parent / "scene_role_lock.json"),
"manifest_hashes": {
n: sha(output / n) for n in ["dev_manifest.json", "holdout_manifest.json"]
},
}
write_json(output.parent / "qualified_frame_role_lock.json", role_lock)
write_json(
output / "download_cost.json",
{
"network_bytes": sum(r.transferred for r in readers),
"budget_accounted_bytes": shared_budget[0],
"full_archives_downloaded": False,
},
)
return manifest
finally:
for obj in archives + readers + sessions:
obj.close()