feat: initial commit - massive-cloud point cloud processing system

scripts/assembly_nearest_gpu.py

@@ -0,0 +1,143 @@
+#!/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()