268 lines
10 KiB
Python
268 lines
10 KiB
Python
#!/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()
|