#!/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()