144 lines
5.1 KiB
Python
144 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
服务端 GPU：PyTorch3D point_face_distance（CUDA）计算点到各零件 STL 三角面的最近距离，
|
||
与 three-mesh-bvh 语义一致（点到三角形）。无 CUDA / 无依赖时退出码 2，由 Node 回退 CPU。
|
||
|
||
依赖（需与 CUDA 版本匹配自行安装）:
|
||
pip install torch trimesh numpy
|
||
pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
|
||
# 或官方 wheel: https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--points", required=True, help="raw float32 little-endian, N*3 bytes")
|
||
ap.add_argument("--manifest-dir", required=True, help="含 manifest.json 与各 .stl")
|
||
ap.add_argument("--out-dir", required=True, help="写入 assignments.int32 / distances.f32 / result.json")
|
||
ap.add_argument("--chunk", type=int, default=65536, help="每批点数，显存不足时调小")
|
||
args = ap.parse_args()
|
||
|
||
try:
|
||
import numpy as np
|
||
import torch
|
||
import trimesh
|
||
from pytorch3d.loss.point_mesh_distance import point_face_distance
|
||
except ImportError as e: # noqa: PERF203
|
||
print(json.dumps({"ok": False, "error": f"import: {e}", "fallback": True}), flush=True)
|
||
sys.exit(2)
|
||
|
||
if not torch.cuda.is_available():
|
||
print(
|
||
json.dumps({"ok": False, "error": "no CUDA device", "fallback": True}),
|
||
flush=True,
|
||
)
|
||
sys.exit(2)
|
||
|
||
device = torch.device("cuda:0")
|
||
pts_path = args.points
|
||
man_dir = args.manifest_dir.rstrip(os.sep)
|
||
out_dir = args.out_dir
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
with open(pts_path, "rb") as f:
|
||
raw = f.read()
|
||
|
||
n_float = len(raw) // 4
|
||
if n_float % 3 != 0:
|
||
print(json.dumps({"ok": False, "error": "points file size not multiple of 12"}), flush=True)
|
||
sys.exit(1)
|
||
pts_np = np.frombuffer(raw, dtype=np.float32).reshape(-1, 3)
|
||
n_pt = int(pts_np.shape[0])
|
||
|
||
man_path = os.path.join(man_dir, "manifest.json")
|
||
with open(man_path, encoding="utf-8") as f:
|
||
man = json.load(f)
|
||
parts = man.get("parts") or []
|
||
if not parts:
|
||
print(json.dumps({"ok": False, "error": "empty manifest.parts"}), flush=True)
|
||
sys.exit(1)
|
||
|
||
stl_infos: list[tuple[torch.Tensor, torch.Tensor]] = []
|
||
for p in parts:
|
||
rel = p.get("stlRelative") or p.get("stl_relative")
|
||
if not rel:
|
||
print(json.dumps({"ok": False, "error": "part missing stlRelative"}), flush=True)
|
||
sys.exit(1)
|
||
stl_path = os.path.join(man_dir, rel)
|
||
if not os.path.isfile(stl_path):
|
||
print(json.dumps({"ok": False, "error": f"missing {stl_path}"}), flush=True)
|
||
sys.exit(1)
|
||
mesh = trimesh.load(stl_path, force="mesh")
|
||
if hasattr(mesh, "geometry"):
|
||
mesh = trimesh.util.concatenate(tuple(mesh.geometry.values()))
|
||
verts = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
|
||
faces = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
|
||
if faces.numel() == 0:
|
||
print(json.dumps({"ok": False, "error": f"empty mesh {rel}"}), flush=True)
|
||
sys.exit(1)
|
||
stl_infos.append((verts, faces))
|
||
|
||
n_parts = len(stl_infos)
|
||
chunk = max(1024, int(args.chunk))
|
||
|
||
# 增量 argmin，避免 (N × 零件数) 距离矩阵占满显存
|
||
best_d = torch.full((n_pt,), float("inf"), dtype=torch.float32, device=device)
|
||
best_i = torch.zeros((n_pt,), dtype=torch.int64, device=device)
|
||
p_first = torch.tensor([0], dtype=torch.int64, device=device)
|
||
t_first = torch.tensor([0], dtype=torch.int64, device=device)
|
||
|
||
for pi in range(n_parts):
|
||
verts, faces = stl_infos[pi]
|
||
tris = verts[faces] # (F, 3, 3)
|
||
for s in range(0, n_pt, chunk):
|
||
end = min(s + chunk, n_pt)
|
||
block = torch.tensor(pts_np[s:end], dtype=torch.float32, device=device)
|
||
mp = end - s
|
||
d_sq = point_face_distance(
|
||
block,
|
||
p_first,
|
||
tris,
|
||
t_first,
|
||
mp,
|
||
5e-3,
|
||
)
|
||
d = torch.sqrt(torch.clamp(d_sq, min=0.0))
|
||
sl = slice(s, end)
|
||
improve = d < best_d[sl]
|
||
best_d[sl] = torch.where(improve, d, best_d[sl])
|
||
best_i[sl] = torch.where(
|
||
improve,
|
||
torch.full((mp,), pi, dtype=torch.int64, device=device),
|
||
best_i[sl],
|
||
)
|
||
|
||
assign = best_i.to(torch.int32).cpu().numpy()
|
||
min_d = best_d.to(torch.float32).cpu().numpy()
|
||
|
||
a_path = os.path.join(out_dir, "assignments.int32")
|
||
d_path = os.path.join(out_dir, "distances.f32")
|
||
assign.tofile(a_path)
|
||
min_d.tofile(d_path)
|
||
|
||
meta = {
|
||
"ok": True,
|
||
"n": int(n_pt),
|
||
"parts": n_parts,
|
||
"assignmentsFile": a_path,
|
||
"distancesFile": d_path,
|
||
"device": str(device),
|
||
}
|
||
with open(os.path.join(out_dir, "result.json"), "w", encoding="utf-8") as f:
|
||
json.dump(meta, f)
|
||
print(json.dumps(meta), flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|