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

scripts/deviation_heatmap_gpu.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+"""
+GPU 偏差热力：PyTorch CUDA 计算网格顶点到最近点云点的距离并着色。
+支持千万～亿级点云：点云不全部 load 进显存，分块流式处理。
+无 CUDA / 无依赖时退出码 2，由 Node 回退 CPU。
+
+用法:
+  python3 deviation_heatmap_gpu.py \
+    --cad /path/to/model.obj \
+    --points /tmp/points.bin \
+    --out-dir /tmp/dev-heatmap \
+    [--chunk-v 2048] [--chunk-p 100000]
+
+输入:
+  --cad       STL 或 OBJ 文件路径
+  --points    采样点 float32 little-endian，N*3 个 float（x y z 交替）
+
+输出（写入 --out-dir）:
+  colors.bin     float32 little-endian，V*3（每顶点 RGB [0,1]）
+  result.json    { ok, vertexCount, dispMin, dispMax }
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import sys
+import time
+
+
+def normalize_like_viewer(pts):
+    """居中 + 缩放到约 10 单位"""
+    mn = pts.min(dim=0).values
+    mx = pts.max(dim=0).values
+    center = (mn + mx) / 2
+    max_dim = (mx - mn).max().item()
+    scale = 10.0 / max_dim if max_dim > 0 else 1.0
+    return (pts - center) * scale, scale
+
+
+def sequential_ramp_rgb(t):
+    """顺序色标 t∈[0,1]：蓝→绿→黄→红"""
+    import torch
+    stops_t = [0.0, 0.33, 0.66, 1.0]
+    stops_rgb = [
+        (0x21 / 255, 0x66 / 255, 0xAC / 255),
+        (0x4D / 255, 0xAF / 255, 0x4A / 255),
+        (0xFE / 255, 0xE0 / 255, 0x8B / 255),
+        (0xD7 / 255, 0x30 / 255, 0x27 / 255),
+    ]
+    t = t.clamp(0.0, 1.0)
+    r = torch.zeros_like(t)
+    g = torch.zeros_like(t)
+    b = torch.zeros_like(t)
+    for i in range(len(stops_t) - 1):
+        t0, t1 = stops_t[i], stops_t[i + 1]
+        r0, g0, b0 = stops_rgb[i]
+        r1, g1, b1 = stops_rgb[i + 1]
+        mask = (t >= t0) & (t <= t1) if i < len(stops_t) - 2 else (t >= t0)
+        s = ((t - t0) / (t1 - t0)).clamp(0, 1)
+        r = torch.where(mask, r0 + (r1 - r0) * s, r)
+        g = torch.where(mask, g0 + (g1 - g0) * s, g)
+        b = torch.where(mask, b0 + (b1 - b0) * s, b)
+    return torch.stack([r, g, b], dim=-1)
+
+
+def auto_tune_chunks(vram_bytes: int, n_verts: int, n_pts: int):
+    """根据可用显存自动决定分块大小，预留 1GB 给网格/点云/开销"""
+    usable = max(vram_bytes - 1024 * 1024 * 1024, 512 * 1024 * 1024)
+    # cdist 输出 float32: chunk_v * chunk_p * 4 bytes
+    # 目标：每个子矩阵 ≤ usable 的 50%
+    budget = usable * 0.5
+    chunk_v = min(4096, n_verts)
+    chunk_p = int(budget / (chunk_v * 4))
+    chunk_p = max(10000, min(chunk_p, n_pts))
+    return chunk_v, chunk_p
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--cad", required=True, help="STL 或 OBJ 文件")
+    ap.add_argument("--points", required=True, help="float32 little-endian N*3")
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--chunk-v", type=int, default=0, help="顶点分块，0=自动")
+    ap.add_argument("--chunk-p", type=int, default=0, help="点云分块，0=自动")
+    ap.add_argument("--smooth", type=int, default=40,
+                    help="拉普拉斯平滑迭代次数，0=不平滑，越大区域越均匀")
+    ap.add_argument("--smooth-alpha", type=float, default=0.7,
+                    help="平滑权重 0~1，越大越平滑")
+    args = ap.parse_args()
+
+    try:
+        import numpy as np
+        import torch
+        import trimesh
+    except ImportError as e:
+        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")
+    vram = torch.cuda.get_device_properties(device).total_memory
+    print(f"[DeviationGPU] VRAM={vram / 1e9:.1f}GB device={torch.cuda.get_device_name(device)}", flush=True)
+
+    t0 = time.time()
+
+    # ── 加载 CAD 网格 ──
+    cad_path = args.cad
+    ext = os.path.splitext(cad_path)[1].lower()
+    mesh = trimesh.load(cad_path, force="mesh")
+    if hasattr(mesh, "geometry"):
+        mesh = trimesh.util.concatenate(tuple(mesh.geometry.values()))
+    verts_np = np.array(mesh.vertices, dtype=np.float32)
+    n_verts = verts_np.shape[0]
+    if n_verts == 0:
+        print(json.dumps({"ok": False, "error": "empty mesh"}), flush=True)
+        sys.exit(1)
+    print(f"[DeviationGPU] mesh {ext} verts={n_verts:,}", flush=True)
+
+    # ── 加载点云（mmap，不占全部内存）──
+    pts_path = args.points
+    file_size = os.path.getsize(pts_path)
+    n_floats = file_size // 4
+    if n_floats % 3 != 0:
+        print(json.dumps({"ok": False, "error": "points size not multiple of 12"}), flush=True)
+        sys.exit(1)
+    n_pts = n_floats // 3
+    pts_mmap = np.memmap(pts_path, dtype=np.float32, mode="r", shape=(n_pts, 3))
+    print(f"[DeviationGPU] points={n_pts:,} ({file_size / 1e9:.2f}GB)", flush=True)
+
+    faces_np = np.array(mesh.faces, dtype=np.int64)
+    n_faces = faces_np.shape[0]
+    print(f"[DeviationGPU] faces={n_faces:,} face_verts={n_faces * 3:,} (Three.js layout)", flush=True)
+
+    # ── 归一化 mesh（放 GPU）──
+    verts_t = torch.tensor(verts_np, dtype=torch.float32, device=device)
+    verts_t, _ = normalize_like_viewer(verts_t)
+
+    # ── 点云归一化参数（在 CPU 上扫一遍 mmap 求 bbox）──
+    SCAN_CAP = min(n_pts, 2_000_000)
+    scan_idx = np.linspace(0, n_pts - 1, SCAN_CAP, dtype=int)
+    scan_pts = pts_mmap[scan_idx]
+    pc_min = scan_pts.min(axis=0)
+    pc_max = scan_pts.max(axis=0)
+    pc_center = (pc_min + pc_max) / 2
+    pc_max_dim = float((pc_max - pc_min).max())
+    pc_scale = 10.0 / pc_max_dim if pc_max_dim > 0 else 1.0
+
+    # ── 分块大小 ──
+    chunk_v = args.chunk_v if args.chunk_v > 0 else 0
+    chunk_p = args.chunk_p if args.chunk_p > 0 else 0
+    if chunk_v == 0 or chunk_p == 0:
+        auto_v, auto_p = auto_tune_chunks(vram, n_verts, n_pts)
+        if chunk_v == 0:
+            chunk_v = auto_v
+        if chunk_p == 0:
+            chunk_p = auto_p
+    sub_mem = chunk_v * chunk_p * 4 / 1024 / 1024
+    total_iters = math.ceil(n_verts / chunk_v) * math.ceil(n_pts / chunk_p)
+    print(f"[DeviationGPU] chunk_v={chunk_v:,} chunk_p={chunk_p:,} sub-matrix≈{sub_mem:.0f}MB iters={total_iters:,}", flush=True)
+
+    # ── 双重分块 KNN ──
+    best_dist = torch.full((n_verts,), float("inf"), dtype=torch.float32, device=device)
+    done_iters = 0
+    last_log = time.time()
+
+    for sv in range(0, n_verts, chunk_v):
+        ev = min(sv + chunk_v, n_verts)
+        block_v = verts_t[sv:ev]  # (Bv, 3)
+
+        for sp in range(0, n_pts, chunk_p):
+            ep = min(sp + chunk_p, n_pts)
+            # 从 mmap 读一块点云 → GPU，归一化
+            block_p_np = np.array(pts_mmap[sp:ep])  # copy from mmap
+            block_p = torch.tensor(block_p_np, dtype=torch.float32, device=device)
+            block_p = (block_p - torch.tensor(pc_center, dtype=torch.float32, device=device)) * pc_scale
+
+            d = torch.cdist(block_v, block_p)  # (Bv, Bp)
+            min_d, _ = d.min(dim=1)
+            best_dist[sv:ev] = torch.minimum(best_dist[sv:ev], min_d)
+            done_iters += 1
+
+        now = time.time()
+        if now - last_log >= 2.0 or sv == 0:
+            pct = done_iters / total_iters * 100
+            elapsed = now - t0
+            eta = elapsed / max(done_iters, 1) * (total_iters - done_iters)
+            print(f"[DeviationGPU] {pct:.0f}% verts={sv:,}/{n_verts:,} elapsed={elapsed:.1f}s ETA={eta:.0f}s", flush=True)
+            last_log = now
+
+    best_dist_cpu = best_dist.cpu()
+    elapsed = time.time() - t0
+    print(f"[DeviationGPU] KNN done in {elapsed:.1f}s", flush=True)
+
+    # ── 拉普拉斯平滑：相邻顶点距离互相平均，消除逐顶点噪声，形成区域级着色 ──
+    smooth_iters = args.smooth
+    smooth_alpha = args.smooth_alpha
+    if smooth_iters > 0:
+        t_smooth = time.time()
+        edges = mesh.edges_unique  # (E, 2) 唯一边
+        e0 = edges[:, 0]
+        e1 = edges[:, 1]
+
+        # 每个顶点的邻居数（度）
+        degree = np.zeros(n_verts, dtype=np.float32)
+        np.add.at(degree, e0, 1)
+        np.add.at(degree, e1, 1)
+        degree[degree == 0] = 1.0
+
+        dist = best_dist_cpu.numpy().copy()
+        alpha = smooth_alpha
+        for it in range(smooth_iters):
+            # 对每个顶点，累加所有邻居的距离值
+            neighbor_sum = np.zeros(n_verts, dtype=np.float64)
+            np.add.at(neighbor_sum, e0, dist[e1])
+            np.add.at(neighbor_sum, e1, dist[e0])
+            neighbor_avg = (neighbor_sum / degree).astype(np.float32)
+            dist = alpha * neighbor_avg + (1 - alpha) * dist
+
+        best_dist_cpu = torch.from_numpy(dist)
+        print(f"[DeviationGPU] smoothed {smooth_iters} iters (alpha={smooth_alpha}) in {time.time() - t_smooth:.1f}s", flush=True)
+
+    # ── 色标范围（4%–96% 分位数）──
+    sorted_d, _ = best_dist_cpu.sort()
+    m = sorted_d.numel()
+    lo = sorted_d[max(0, int(m * 0.04))].item()
+    hi = sorted_d[min(m - 1, int(math.ceil(m * 0.96)) - 1)].item()
+    if hi <= lo:
+        hi = lo + 1e-9
+    disp_min = float(lo)
+    disp_max = float(hi)
+
+    # ── 展开为 per-face-vertex（与 Three.js OBJLoader 非索引布局一致）──
+    face_vert_idx = torch.from_numpy(faces_np.flatten().astype(np.int64))
+    best_dist_face = best_dist_cpu[face_vert_idx]  # (F*3,)
+    n_face_verts = best_dist_face.numel()
+    print(f"[DeviationGPU] expanded unique→face verts: {n_verts:,} → {n_face_verts:,}", flush=True)
+
+    # ── 着色 ──
+    t_norm = ((best_dist_face - disp_min) / (disp_max - disp_min)).clamp(0, 1)
+    colors = sequential_ramp_rgb(t_norm)  # (F*3, 3)
+
+    # ── 写出 ──
+    os.makedirs(args.out_dir, exist_ok=True)
+    colors_np = colors.numpy().astype(np.float32)
+    colors_path = os.path.join(args.out_dir, "colors.bin")
+    colors_np.tofile(colors_path)
+
+    meta = {
+        "ok": True,
+        "vertexCount": int(n_face_verts),
+        "dispMin": disp_min,
+        "dispMax": disp_max,
+        "pointsUsed": int(n_pts),
+        "elapsedSec": round(time.time() - t0, 1),
+    }
+    with open(os.path.join(args.out_dir, "result.json"), "w", encoding="utf-8") as f:
+        json.dump(meta, f)
+    print(json.dumps(meta), flush=True)
+
+
+if __name__ == "__main__":
+    main()