File size: 6,251 Bytes
0e4e948 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """
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)
# 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()
|