File size: 4,483 Bytes
30857ce | 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 | #!/usr/bin/env python3
"""
CLI script: reconstruct a mesh from a point-cloud file using NKSR.
Usage
-----
python reconstruct.py input.ply output.ply --detail 1.0 --mise-iter 1
python reconstruct.py input.ply output.ply --chunk-size 50.0 --no-normals
"""
import argparse
import sys
from pathlib import Path
import torch
# Allow running from repo root without installing
sys.path.insert(0, str(Path(__file__).parent.parent))
from nksr_wrapper import NKSRMeshReconstructor, load_point_cloud, save_mesh
def main() -> None:
parser = argparse.ArgumentParser(
description="NKSR point-cloud → mesh reconstruction"
)
parser.add_argument("input", type=Path, help="Input PLY or PCD file")
parser.add_argument("output", type=Path, help="Output mesh file (PLY/OBJ/GLB)")
parser.add_argument("--device", default="cuda:0", help="PyTorch device")
parser.add_argument("--config", default="ks", help="NKSR model config (ks/snet/snet-wonormal)")
parser.add_argument("--detail", type=float, default=1.0, help="Detail level 0.0-1.0")
parser.add_argument("--voxel-size", type=float, default=None, help="Override voxel size")
parser.add_argument("--chunk-size", type=float, default=-1.0, help="Chunk size for large scenes")
parser.add_argument("--mise-iter", type=int, default=1, help="MISE iterations")
parser.add_argument("--no-normals", action="store_true", help="Ignore normals in file; estimate them")
parser.add_argument("--estimate-normals", action="store_true", help="Estimate normals if file lacks them")
parser.add_argument("--sensor", type=Path, default=None, help="Optional NPY file with sensor positions")
parser.add_argument("--colors", type=Path, default=None, help="Optional NPY file with per-point RGB colors")
parser.add_argument("--solver-iter", type=int, default=2000, help="PCG solver max iterations")
parser.add_argument("--solver-tol", type=float, default=1e-5, help="PCG solver tolerance")
parser.add_argument("--verbose", action="store_true", help="Print extra progress info")
args = parser.parse_args()
if not args.input.exists():
parser.error(f"Input file not found: {args.input}")
# ---- load point cloud -----------------------------------------------
print(f"Loading point cloud from {args.input} ...")
points, normals = load_point_cloud(
args.input,
estimate_normals=args.estimate_normals or args.no_normals,
)
print(f" Loaded {len(points)} points")
if normals is not None:
print(f" Normals present: {normals.shape}")
elif not args.no_normals:
print(" No normals found in file — will estimate on-the-fly")
if args.no_normals:
normals = None
print(" --no-normals set: normals will be estimated")
# ---- optional extras ------------------------------------------------
sensor = None
if args.sensor:
import numpy as np
sensor = np.load(args.sensor)
print(f" Sensor positions loaded: {sensor.shape}")
colors = None
if args.colors:
import numpy as np
colors = np.load(args.colors)
print(f" Per-point colors loaded: {colors.shape}")
# ---- reconstruct ----------------------------------------------------
print("\nInitialising NKSR reconstructor ...")
if not torch.cuda.is_available() and args.device.startswith("cuda"):
print("WARNING: CUDA not available, falling back to CPU (very slow)")
args.device = "cpu"
recon = NKSRMeshReconstructor(
device=args.device,
config=args.config,
chunk_tmp_device="cpu" if args.chunk_size > 0 else None,
)
print("Reconstructing mesh ...")
mesh = recon.reconstruct(
points=points,
normals=normals,
sensor_positions=sensor,
colors=colors,
detail_level=args.detail,
voxel_size=args.voxel_size,
chunk_size=args.chunk_size,
mise_iter=args.mise_iter,
solver_max_iter=args.solver_iter,
solver_tol=args.solver_tol,
)
# ---- save -----------------------------------------------------------
print(f"\nSaving mesh to {args.output} ...")
save_mesh(args.output, mesh.vertices, mesh.faces, mesh.vertex_colors)
print(f" Vertices: {len(mesh.vertices):,} | Faces: {len(mesh.faces):,}")
if mesh.vertex_colors is not None:
print(" Vertex colors included")
print("Done.")
if __name__ == "__main__":
main()
|