File size: 4,023 Bytes
eafbe80 | 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 | #!/usr/bin/env python3
"""
Lightweight DATASET sanity check before heavy eval (no torch / no train.py).
Verifies jsons + frames dirs and that pose keys exist for the replay frame range.
Exit 0 if OK, 1 otherwise; prints absolute paths and first failure reason.
"""
from __future__ import annotations
import argparse
import json
import os
import sys
def _abs(p: str) -> str:
return os.path.abspath(os.path.expanduser(p))
def _poses_dict(data: dict) -> dict:
if "CineCameraActor" in data:
return data["CineCameraActor"]
return data if isinstance(data, dict) else {}
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--dataset", required=True)
ap.add_argument("--video", required=True)
ap.add_argument("--start_frame", type=int, default=0)
ap.add_argument("--num_chunks", type=int, default=1)
ap.add_argument("--chunk_frames", type=int, default=81)
args = ap.parse_args()
ds = _abs(args.dataset)
jsons = os.path.join(ds, "jsons")
frames_root = os.path.join(ds, "frames")
vn = str(args.video).replace(".mp4", "").replace(".avi", "").strip()
jpath = os.path.join(jsons, f"{vn}.json")
fdir = os.path.join(frames_root, vn)
print(f"[check_dataset_gt] DATASET={ds}")
print(f"[check_dataset_gt] json={jpath}")
print(f"[check_dataset_gt] frames_dir={fdir}")
if not os.path.isdir(ds):
print("[check_dataset_gt] FAIL: DATASET is not a directory", file=sys.stderr)
return 1
if not os.path.isdir(jsons):
print("[check_dataset_gt] FAIL: missing jsons/", file=sys.stderr)
return 1
if not os.path.isfile(jpath):
print("[check_dataset_gt] FAIL: missing video json", file=sys.stderr)
return 1
if not os.path.isdir(fdir):
print("[check_dataset_gt] WARN: frames subdir missing (PNG compare may fail)", file=sys.stderr)
try:
with open(jpath, "r", encoding="utf-8") as f:
data = json.load(f)
except Exception as e:
print(f"[check_dataset_gt] FAIL: cannot read json: {e}", file=sys.stderr)
return 1
poses = _poses_dict(data)
need = []
for ch in range(args.num_chunks):
seg = args.start_frame + ch * args.chunk_frames
for i in range(args.chunk_frames):
need.append(seg + i)
missing_pose = []
for fi in need:
k = str(fi)
if k not in poses:
missing_pose.append(fi)
if len(missing_pose) >= 5:
break
if missing_pose:
print(
f"[check_dataset_gt] FAIL: missing pose keys for frames (showing up to 5): {missing_pose}",
file=sys.stderr,
)
def _knum(x):
try:
return int(x)
except (TypeError, ValueError):
return 0
sample_keys = sorted(poses.keys(), key=_knum)[:8]
print(f"[check_dataset_gt] sample pose keys: {sample_keys}", file=sys.stderr)
return 1
# Optional: same 12-dim path as replay (numpy via fov_retrieval; no torch)
try:
_here = os.path.dirname(os.path.abspath(__file__))
# .../eval/v2/basic -> repo root
_repo = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(_here))))
if _repo not in sys.path:
sys.path.insert(0, _repo)
from src.model_training.fov_retrieval import load_camera_pose, pose_to_rt
for fi in need[: min(3, len(need))]:
pose = load_camera_pose(jpath, fi)
rt = pose_to_rt(pose, constrain_to_xy=True) if pose else None
if rt is None or len(rt) < 12:
print(f"[check_dataset_gt] FAIL: pose_to_rt None for frame {fi}", file=sys.stderr)
return 1
except Exception as e:
print(f"[check_dataset_gt] WARN: RT parse spot-check skipped: {e}", file=sys.stderr)
print("[check_dataset_gt] OK: paths and pose keys cover the replay range.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|