""" Single-mesh pipeline (SQ fitting only): 1. load + normalize mesh 2. build adaptive block map (or load existing) 3. compute SDF volume (or load existing) 4. run SDFAdaptiveFitter 5. save results Example: python scripts/run_pipeline.py \ --mesh_path /path/to/watertight_mesh.ply \ --output_root /path/to/output \ --complex_levels 16,32 \ --max_batches 50 --max_cands 5 --target_coverage 0.9 --max_iter 200 """ import os import sys import argparse import json import numpy as np import open3d as o3d import trimesh from hitops.adaptive.adaptive_res_block import build_adaptive_block_map, save_adaptive_map from hitops.sdf.build_adaptive_sdf import ( build_sdf_volume, save_sdf_volume, load_sdf_volume, sdf_volume_stats, ) from hitops.sqfit.sq_fit import SDFAdaptiveFitter # ----------------------------------------------------------------------- # Utility functions # ----------------------------------------------------------------------- def normalize_to_unit_cube_safe(mesh, half=0.5, margin=1e-6): v = np.asarray(mesh.vertices).copy() f = np.asarray(mesh.triangles).copy() if v.size == 0 or f.size == 0: raise ValueError("Empty mesh.") if not np.isfinite(v).all(): raise ValueError("Mesh has NaN/Inf vertices.") center = (v.min(0) + v.max(0)) * 0.5 max_ext = float((v.max(0) - v.min(0)).max()) if max_ext <= 0: raise ValueError(f"Invalid extent: {max_ext}") v = (v - center) * ((half - margin) * 2.0 / max_ext) out = o3d.geometry.TriangleMesh() out.vertices = o3d.utility.Vector3dVector(v) out.triangles = o3d.utility.Vector3iVector(f) return out def _load_adaptive_input(path: str) -> dict: """Load an adaptive map from an npz file or an npy dict.""" loaded = np.load(path, allow_pickle=True) if isinstance(loaded, np.lib.npyio.NpzFile): if "level_map" in loaded and "block_size_map" in loaded: return { "level_map": loaded["level_map"].astype(np.int32), "block_size_map": loaded["block_size_map"].astype(np.int32), } raise ValueError(f"NPZ missing required keys: {path}") obj = loaded.item() if "level_map" in obj and "block_size_map" in obj: return { "level_map": np.asarray(obj["level_map"], dtype=np.int32), "block_size_map": np.asarray(obj["block_size_map"], dtype=np.int32), } if "map" in obj: lm = np.asarray(obj["map"], dtype=np.int32) print("[Warn] Legacy 'map' key, using fallback block_size_map=32.") return {"level_map": lm, "block_size_map": np.where(lm > 0, 32, 0).astype(np.int32)} raise ValueError(f"Cannot parse input file: {path}") # ----------------------------------------------------------------------- # Main pipeline # ----------------------------------------------------------------------- def run_pipeline( mesh_path: str, output_root: str, b_list: list, r_max: int, level_scheme: str, clip_percentile: float, complex_levels: list, max_batches: int, max_cands: int, target_coverage: float, max_iter: int, cost_threshold: float, max_total_sqs: int, map_npz_path: str = None, # existing adaptive map sdf_npz_path: str = None, # existing SDF volume ): if not os.path.exists(mesh_path): raise FileNotFoundError(f"Mesh not found: {mesh_path}") mesh_name = os.path.splitext(os.path.basename(mesh_path))[0] out_dir = os.path.join(output_root, mesh_name) map_dir = os.path.join(out_dir, "adaptive_map") sdf_dir = os.path.join(out_dir, "sdf_vol") sq_dir = os.path.join(out_dir, "sq_fit_v20") for d in [map_dir, sdf_dir, sq_dir]: os.makedirs(d, exist_ok=True) # ---- Step 1: load mesh ---- print(f"\n[1/4] Loading mesh: {mesh_path}") mesh = o3d.io.read_triangle_mesh(mesh_path) if len(mesh.vertices) == 0: raise ValueError("Empty mesh.") print(f" Vertices: {len(mesh.vertices)}, Triangles: {len(mesh.triangles)}") # ---- Step 2: normalize ---- print("[2/4] Normalizing to [-0.5, 0.5]^3") mesh = normalize_to_unit_cube_safe(mesh) # ---- Step 3: Adaptive block map ---- if map_npz_path and os.path.exists(map_npz_path): print(f"[3a/4] Loading pre-built adaptive map: {map_npz_path}") adaptive_map = _load_adaptive_input(map_npz_path) saved_map = map_npz_path else: print("[3a/4] Building adaptive block map") adaptive_map = build_adaptive_block_map( mesh=mesh, B_list=b_list, R_max=r_max, level_scheme=level_scheme, clip_percentile=clip_percentile, ) saved_map = os.path.join(map_dir, f"{mesh_name}_adaptive_map.npz") save_adaptive_map(adaptive_map, saved_map) # ---- Step 3b: SDF volume ---- R = adaptive_map["level_map"].shape[0] if sdf_npz_path and os.path.exists(sdf_npz_path): print(f"[3b/4] Loading pre-built SDF volume: {sdf_npz_path}") sdf_vol, truncation = load_sdf_volume(sdf_npz_path) saved_sdf = sdf_npz_path else: print(f"[3b/4] Computing SDF volume at {R}³ resolution...") sdf_vol, truncation = build_sdf_volume(mesh, resolution=R) saved_sdf = os.path.join(sdf_dir, f"{mesh_name}_sdf_vol.npz") save_sdf_volume(sdf_vol, truncation, saved_sdf) print(" SDF stats:") sdf_volume_stats(sdf_vol, truncation) # Report statistics on complex (high-resolution) regions lmap = adaptive_map["level_map"] total_occ = int(np.sum(lmap > 0)) comp_cnt = int(np.sum(np.isin(lmap, complex_levels))) print(f"\n complex_levels={complex_levels}: {comp_cnt}/{total_occ} ({comp_cnt/max(total_occ,1)*100:.1f}%)") # ---- Step 4: v20 SQ fitting ---- print(f"\n[4/4] Running v20 SDFAdaptiveFitter (cost_threshold={cost_threshold})") fitter = SDFAdaptiveFitter( adaptive_map=adaptive_map, sdf_vol=sdf_vol, truncation=truncation, complex_levels=tuple(complex_levels), ) fitter.run( max_batches=max_batches, max_cands=max_cands, target_coverage=target_coverage, max_iter=max_iter, cost_threshold=cost_threshold, max_total_sqs=max_total_sqs, ) fitter.save_results(sq_dir) # ---- Summary ---- summary = { "mesh_path": mesh_path, "mesh_name": mesh_name, "adaptive_map": saved_map, "sdf_vol": saved_sdf, "sq_dir": sq_dir, "n_sq": len(fitter.all_sqs_params), "complex_levels": list(complex_levels), "complex_voxel_ratio": f"{comp_cnt/max(total_occ,1)*100:.1f}%", "params": { "B_list": b_list, "R_max": r_max, "level_scheme": level_scheme, "clip_percentile": clip_percentile, "complex_levels": list(complex_levels), "max_batches": max_batches, "max_cands": max_cands, "target_coverage": target_coverage, "max_iter": max_iter, "cost_threshold": cost_threshold, }, } summary_path = os.path.join(out_dir, "run_summary_v20.json") with open(summary_path, "w") as f: json.dump(summary, f, indent=2) print(f"\n=== Done ===") print(f"Adaptive map : {saved_map}") print(f"SDF volume : {saved_sdf}") print(f"SQ outputs : {sq_dir} ({len(fitter.all_sqs_params)} SQs)") print(f"Summary : {summary_path}") # ----------------------------------------------------------------------- # CLI # ----------------------------------------------------------------------- def parse_int_list(s): return [int(x.strip()) for x in s.split(",") if x.strip()] def main(): parser = argparse.ArgumentParser(description="v20 adaptive SDF + SQ fitting pipeline") parser.add_argument("--mesh_path", type=str, required=True) parser.add_argument("--output_root", type=str, required=True, help="output root; each mesh is written to //") parser.add_argument("--b_list", type=str, default="16,32,64") parser.add_argument("--r_max", type=int, default=32) parser.add_argument("--level_scheme", type=str, default="quantile", choices=["quantile", "fixed"]) parser.add_argument("--clip_percentile", type=float, default=99.5) parser.add_argument("--complex_levels", type=str, default="16,32") parser.add_argument("--max_batches", type=int, default=50) parser.add_argument("--max_cands", type=int, default=5) parser.add_argument("--target_coverage", type=float, default=0.9) parser.add_argument("--max_iter", type=int, default=200) parser.add_argument("--cost_threshold", type=float, default=0.05, help="max SDF fitting error; SQs above this are discarded") parser.add_argument("--max_total_sqs", type=int, default=50) parser.add_argument("--map_npz_path", type=str, default=None, help="path to an existing adaptive map .npz; skips map construction") parser.add_argument("--sdf_npz_path", type=str, default=None, help="path to an existing SDF volume .npz; skips SDF computation") args = parser.parse_args() run_pipeline( mesh_path=args.mesh_path, output_root=args.output_root, b_list=parse_int_list(args.b_list), r_max=args.r_max, level_scheme=args.level_scheme, clip_percentile=args.clip_percentile, complex_levels=parse_int_list(args.complex_levels), max_batches=args.max_batches, max_cands=args.max_cands, target_coverage=args.target_coverage, max_iter=args.max_iter, cost_threshold=args.cost_threshold, max_total_sqs=args.max_total_sqs, map_npz_path=args.map_npz_path, sdf_npz_path=args.sdf_npz_path, ) if __name__ == "__main__": main()