File size: 4,823 Bytes
4c3d957 a3752bd 4c3d957 a3752bd 4c3d957 | 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 119 120 | #!/usr/bin/env python
"""Align post-hoc (non-geometry-grounded) baseline meshes (Pixal3D / RVG) to the
GT canonical [-0.5,0.5]^3 frame, PRESERVING texture/vertex colour, so the NO-ICP
appeval harness can score them fairly.
For each pred glb: compute cube-rotation + ICP alignment to the GT canon mesh
(align.align_mesh on a geometry copy), then apply that same transform to the
FULL textured mesh (apply_alignment only touches .vertices) and export.
Usage:
python align_baselines.py --pred-dir DIR --gt-dir DIR --out DIR \
--selection SEL.json --limit 30
"""
from __future__ import annotations
import argparse, json, sys
from pathlib import Path
import numpy as np
import trimesh
sys.path.insert(0, str(Path(__file__).resolve().parent))
from faithfulness import voxelize_points
from align import align_mesh, apply_alignment
def gt_voxels(gt_mesh: trimesh.Trimesh, n: int = 64, samples: int = 200_000):
"""Full occupancy of the GT canon mesh (already in [-0.5,0.5]^3)."""
pts = gt_mesh.sample(samples)
return voxelize_points(np.asarray(pts), n)
def load_textured(path: Path) -> trimesh.Trimesh:
m = trimesh.load(str(path), force="mesh", process=False)
return m
def geom_copy(m: trimesh.Trimesh) -> trimesh.Trimesh:
return trimesh.Trimesh(np.asarray(m.vertices), np.asarray(m.faces),
process=False)
def resolve_gt(gt_dir: Path, obj: str):
for c in (gt_dir / f"{obj}_canon.glb", gt_dir / obj / "mesh.glb",
gt_dir / f"{obj}.glb"):
if c.exists():
return c
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--pred-dir", type=Path, required=True)
ap.add_argument("--gt-dir", type=Path, required=True)
ap.add_argument("--out", type=Path, required=True)
ap.add_argument("--selection", type=Path, required=True)
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--n", type=int, default=64)
ap.add_argument("--no-icp", dest="icp", action="store_false", default=True,
help="AXIS-ONLY alignment: unit-cube normalization + best-of-24 "
"discrete octahedral orientation (voxel-F1 scored), and SKIP "
"the trimesh ICP R/t/scale refinement. Default keeps ICP on "
"(existing behaviour unchanged).")
ap.add_argument("--shard", type=int, default=0)
ap.add_argument("--nshards", type=int, default=1)
args = ap.parse_args()
args.out.mkdir(parents=True, exist_ok=True)
sel = json.loads(args.selection.read_text())["selections"]
objs = [s["object"] for s in sel]
if args.limit:
objs = objs[:args.limit]
if args.nshards > 1:
objs = objs[args.shard::args.nshards]
report = []
for i, obj in enumerate(objs, 1):
pred_p = args.pred_dir / f"{obj}.glb"
outp = args.out / f"{obj}.glb"
if outp.exists():
print(f"[{i}/{len(objs)}] {obj} SKIP (aligned exists)", flush=True)
continue
if not pred_p.exists():
print(f"[{i}/{len(objs)}] {obj} MISSING pred", flush=True)
report.append({"object": obj, "error": "missing pred"})
continue
gt_p = resolve_gt(args.gt_dir, obj)
if gt_p is None:
print(f"[{i}/{len(objs)}] {obj} MISSING gt", flush=True)
report.append({"object": obj, "error": "missing gt"})
continue
try:
tex = load_textured(pred_p)
gt = trimesh.load(str(gt_p), force="mesh", process=False)
gt_ref = gt_voxels(gt, n=args.n)
geo = geom_copy(tex)
ar = align_mesh(geo, gt_ref, gt, n=args.n, icp=args.icp)
aligned = apply_alignment(tex, ar) # touches only .vertices
aligned.export(str(outp))
report.append({"object": obj, "f1_raw": round(ar.f1_raw, 4),
"f1_icp": round(ar.f1_icp, 4),
"best_f1": round(ar.best_f1(), 4)})
print(f"[{i}/{len(objs)}] {obj} OK best_f1={ar.best_f1():.3f} "
f"(raw={ar.f1_raw:.3f} icp={ar.f1_icp:.3f})", flush=True)
except Exception as e:
import traceback; traceback.print_exc()
report.append({"object": obj, "error": f"{type(e).__name__}: {e}"})
print(f"[{i}/{len(objs)}] {obj} FAIL {e}", flush=True)
(args.out / "align_report.json").write_text(json.dumps(report, indent=2))
ok = [r for r in report if "best_f1" in r]
if ok:
f1s = [r["best_f1"] for r in ok]
low = [r["object"] for r in ok if r["best_f1"] < 0.4]
print(f"\nALIGN DONE n={len(ok)} mean_best_f1={np.mean(f1s):.3f} "
f"min={np.min(f1s):.3f} | low(<0.4): {low}")
if __name__ == "__main__":
main()
|