| """ |
| Batch mesh-mapper: assign mesh faces to fitted superquadrics (reuses existing |
| SQ-fit + curvature atoms). |
| |
| Inputs (all via CLI): |
| --mesh_dir <mesh_dir>/<uid>/full.ply |
| --v4_dir <v4_dir>/<uid>/face_labels.npy (curvature atoms) |
| --sq_root <sq_root>/<uid>/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 <mesh_dir>/<uid>/full.ply") |
| ap.add_argument("--v4_dir", type=str, required=True, |
| help="curvature-seg dir; labels at <v4_dir>/<uid>/face_labels.npy") |
| ap.add_argument("--sq_root", type=str, required=True, |
| help="SQ-fit root; SQs at <sq_root>/<uid>/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/<uid>/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) |
|
|
| |
| 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_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() |
|
|
