""" Batch mesh-mapper: assign mesh faces to fitted superquadrics (reuses existing SQ-fit + curvature atoms). Inputs (all via CLI): --mesh_dir //full.ply --v4_dir //face_labels.npy (curvature atoms) --sq_root //sq_fit_v20/post_sq_*.ply --output_root destination Per UID: 1. load mesh / face_labels / post_sq_*.ply 2. run mesh_mapper.map_with_atoms 3. write face_labels_v8.npy + mesh_mapped_v8.ply + report.json 4. aggregate to batch_summary.csv / .json Supports --skip_existing for resume. """ from __future__ import annotations import os import sys import csv import json import time import glob import argparse import numpy as np import trimesh from tqdm import tqdm from hitops.mapping.mesh_mapper import ( _load_sq_voxel_groups, map_with_atoms, save_results, ) def discover_uids(v4_dir: str) -> list[str]: uids = [] for d in sorted(os.listdir(v4_dir)): full = os.path.join(v4_dir, d) if os.path.isdir(full) and os.path.isfile(os.path.join(full, "face_labels.npy")): uids.append(d) return uids def process_one( uid: str, mesh_dir: str, v4_dir: str, sq_root: str, output_root: str, *, vote_mode="count", vote_tau=0.025, orphan_face_dist=0.04, orphan_atom_frac=0.6, min_atom_size=3, save_per_sq=False, ) -> dict: mesh_path = os.path.join(mesh_dir, uid, "full.ply") fl_path = os.path.join(v4_dir, uid, "face_labels.npy") sq_dir = os.path.join(sq_root, uid, "sq_fit_v20") out_dir = os.path.join(output_root, uid) for p, name in [(mesh_path, "mesh"), (fl_path, "face_labels"), (sq_dir, "sq_dir")]: if not os.path.exists(p): return {"uid": uid, "status": "skipped", "reason": f"missing {name}: {p}"} t0 = time.time() mesh = trimesh.load(mesh_path, force="mesh", process=False) face_labels_v4 = np.load(fl_path).astype(np.int32) sq_names, sq_pts = _load_sq_voxel_groups(sq_dir) if not sq_names: return {"uid": uid, "status": "skipped", "reason": "empty sq_dir"} res = map_with_atoms( mesh, face_labels_v4, sq_pts, vote_mode=vote_mode, vote_tau=vote_tau, orphan_face_dist=orphan_face_dist, orphan_atom_frac=orphan_atom_frac, min_atom_size=min_atom_size, ) save_results(mesh, res, sq_names, out_dir, save_per_sq=save_per_sq) t_total = time.time() - t0 s = { "uid": uid, "status": "ok", "t_total": round(t_total, 2), **res["stats"], } return s def main(): ap = argparse.ArgumentParser(description="HY3D batch runner for mesh_mapper_v8") ap.add_argument("--mesh_dir", type=str, required=True, help="mesh dir; meshes at //full.ply") ap.add_argument("--v4_dir", type=str, required=True, help="curvature-seg dir; labels at //face_labels.npy") ap.add_argument("--sq_root", type=str, required=True, help="SQ-fit root; SQs at //sq_fit_v20/") ap.add_argument("--output_root", type=str, required=True) ap.add_argument("--vote_mode", type=str, default="count", choices=["count", "exp"]) ap.add_argument("--vote_tau", type=float, default=0.025) ap.add_argument("--orphan_face_dist", type=float, default=0.04) ap.add_argument("--orphan_atom_frac", type=float, default=0.6) ap.add_argument("--min_atom_size", type=int, default=3) ap.add_argument("--save_per_sq", action="store_true") ap.add_argument("--skip_existing", action="store_true", help="skip a UID if output//face_labels_v8.npy already exists") args = ap.parse_args() uids = discover_uids(args.v4_dir) print(f"[batch-v8] {len(uids)} UID(s) to process", flush=True) os.makedirs(args.output_root, exist_ok=True) stats: list[dict] = [] for uid in tqdm(uids, desc="mapper_v8", ncols=100): out_dir = os.path.join(args.output_root, uid) if args.skip_existing and os.path.exists(os.path.join(out_dir, "face_labels_v8.npy")): stats.append({"uid": uid, "status": "skipped", "reason": "already_done"}) continue try: s = process_one( uid, args.mesh_dir, args.v4_dir, args.sq_root, args.output_root, vote_mode=args.vote_mode, vote_tau=args.vote_tau, orphan_face_dist=args.orphan_face_dist, orphan_atom_frac=args.orphan_atom_frac, min_atom_size=args.min_atom_size, save_per_sq=args.save_per_sq, ) except Exception as e: import traceback traceback.print_exc() s = {"uid": uid, "status": "failed", "error": str(e)} stats.append(s) # Write incremental JSON summary after each UID. with open(os.path.join(args.output_root, "batch_summary.json"), "w") as f: json.dump(stats, f, indent=2, default=lambda x: float(x) if isinstance(x, np.floating) else int(x)) # csv csv_path = os.path.join(args.output_root, "batch_summary.csv") keys_order = [ "uid", "status", "t_total", "F", "A", "K", "n_face_orphan_NN", "n_v4_orphan_face", "n_atom_orphan", "atom_size_min", "atom_size_median", "atom_size_max", "atom_conf_mean", "atom_conf_p25", "face_coverage", "sqs_used", ] with open(csv_path, "w", newline="") as f: w = csv.DictWriter(f, fieldnames=keys_order, extrasaction="ignore") w.writeheader() for s in stats: w.writerow({k: s.get(k, "") for k in keys_order}) ok = [s for s in stats if s.get("status") == "ok"] print(f"\n[batch-v8] {len(ok)}/{len(stats)} OK | csv={csv_path}") if ok: print(f" t_total: mean={np.mean([s['t_total'] for s in ok]):.2f}s") print(f" face_coverage: mean={np.mean([s['face_coverage'] for s in ok])*100:.2f}%") print(f" sqs_used/K: " f"mean={np.mean([s['sqs_used']/max(s['K'],1) for s in ok])*100:.1f}%") print(f" atom_conf_mean: mean={np.mean([s['atom_conf_mean'] for s in ok]):.3f}") if __name__ == "__main__": main()