70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Check integrity of processed jobs (PLY + poses present)."""
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
REQUIRED_FILES = ["reconstruction.ply", "lingbot_poses.npz"]
|
|
OPTIONAL_FILES = ["model_decimated.ply"]
|
|
|
|
|
|
def check_job(job_id: int, frames_base: str = "/root/cosma-qc-frames") -> dict[str, Any]:
|
|
job_dir = Path(frames_base) / f"job_{job_id}"
|
|
if not job_dir.exists():
|
|
return {"job_id": job_id, "status": "missing", "missing": [], "details": {}}
|
|
|
|
missing = [f for f in REQUIRED_FILES if not (job_dir / f).exists()]
|
|
details: dict[str, Any] = {}
|
|
|
|
ply = job_dir / "reconstruction.ply"
|
|
if ply.exists():
|
|
details["ply_size_gb"] = round(ply.stat().st_size / 1e9, 2)
|
|
|
|
poses = job_dir / "lingbot_poses.npz"
|
|
if poses.exists():
|
|
try:
|
|
import numpy as np
|
|
d = np.load(str(poses))
|
|
n = d["poses"].shape[0] if "poses" in d else 0
|
|
details["n_poses"] = n
|
|
except Exception as e:
|
|
details["poses_error"] = str(e)
|
|
|
|
decimated = job_dir / "model_decimated.ply"
|
|
details["decimated"] = decimated.exists()
|
|
|
|
return {
|
|
"job_id": job_id,
|
|
"status": "ok" if not missing else "incomplete",
|
|
"missing": missing,
|
|
"details": details,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("job_ids", nargs="*", type=int)
|
|
p.add_argument("--frames-base", default="/root/cosma-qc-frames")
|
|
p.add_argument("--all", action="store_true", help="Check all job dirs")
|
|
args = p.parse_args()
|
|
|
|
base = Path(args.frames_base)
|
|
if args.all:
|
|
ids = sorted(
|
|
int(d.name.replace("job_", ""))
|
|
for d in base.iterdir()
|
|
if d.is_dir() and d.name.startswith("job_") and d.name[4:].isdigit()
|
|
)
|
|
else:
|
|
ids = args.job_ids
|
|
|
|
results = [check_job(jid, args.frames_base) for jid in ids]
|
|
print(json.dumps(results, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|