lipan
29f1d45d66
算法更新: delivery v1.2.1、平整度v1.5、测距v2.2、整车配准v1.3 适配器: vehicle-registration 新增、各算法 manifest 和 adapter 更新 后端: taskCadFiles/source 字段迁移、derivedPointClouds 注册、取消任务联动 前端: ElapsedTime 时长显示、两步算法流程、canRun 逻辑修复 文档: 新增5份算法性能分析报告
1594 lines
50 KiB
Python
1594 lines
50 KiB
Python
import argparse
|
||
import posixpath
|
||
import time
|
||
from pathlib import Path
|
||
from zipfile import ZipFile, ZIP_DEFLATED
|
||
from xml.sax.saxutils import escape
|
||
from xml.etree import ElementTree as ET
|
||
import numpy as np
|
||
import torch
|
||
import trimesh
|
||
import matplotlib.pyplot as plt
|
||
|
||
from OCC.Core.STEPControl import STEPControl_Reader
|
||
from OCC.Core.IFSelect import IFSelect_RetDone
|
||
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
|
||
|
||
from OCC.Core.TopExp import TopExp_Explorer
|
||
from OCC.Core.TopAbs import TopAbs_FACE
|
||
from OCC.Core.BRep import BRep_Tool
|
||
|
||
# 兼容不同 pythonocc 版本
|
||
try:
|
||
from OCC.Core.TopoDS import topods
|
||
USE_TOPODS = "topods"
|
||
except ImportError:
|
||
from OCC.Core.TopoDS import TopoDS
|
||
USE_TOPODS = "TopoDS"
|
||
|
||
try:
|
||
import open3d as o3d
|
||
except ModuleNotFoundError:
|
||
o3d = None
|
||
|
||
|
||
# ============================================================
|
||
# 设备与精度
|
||
# ============================================================
|
||
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
dtype = torch.float64
|
||
|
||
DEFAULT_TASK_IDS = [
|
||
11, 13, 17, 20, 25, 26, 27, 31, 34,
|
||
35, 36, 37, 38, 39, 40, 41, 42, 45
|
||
]
|
||
|
||
DEFAULT_CANDIDATE_TRIANGLE_COUNT = 32
|
||
DEFAULT_RESULT_TEMPLATE = Path("excel/测量结果.xlsx")
|
||
|
||
task_distance_data = {}
|
||
|
||
|
||
# ============================================================
|
||
# STEP 读取
|
||
# ============================================================
|
||
|
||
def read_step_file(step_path):
|
||
reader = STEPControl_Reader()
|
||
status = reader.ReadFile(step_path)
|
||
|
||
if status != IFSelect_RetDone:
|
||
raise RuntimeError(f"STEP 文件读取失败: {step_path}")
|
||
|
||
reader.TransferRoots()
|
||
shape = reader.OneShape()
|
||
|
||
print("STEP 文件读取完成")
|
||
|
||
return shape
|
||
|
||
|
||
# ============================================================
|
||
# OCC Shape -> numpy mesh
|
||
# ============================================================
|
||
|
||
def occ_shape_to_numpy_mesh(shape):
|
||
"""
|
||
从 OCC Shape 中提取三角网格。
|
||
|
||
返回:
|
||
vertices_np: [N, 3]
|
||
faces_np: [M, 3]
|
||
"""
|
||
|
||
vertices = []
|
||
faces = []
|
||
|
||
explorer = TopExp_Explorer(shape, TopAbs_FACE)
|
||
|
||
face_count = 0
|
||
triangle_count = 0
|
||
|
||
while explorer.More():
|
||
current_shape = explorer.Current()
|
||
|
||
if USE_TOPODS == "topods":
|
||
face = topods.Face(current_shape)
|
||
else:
|
||
face = TopoDS.Face(current_shape)
|
||
|
||
location = face.Location()
|
||
|
||
triangulation = BRep_Tool.Triangulation(face, location)
|
||
|
||
if triangulation is None:
|
||
explorer.Next()
|
||
continue
|
||
|
||
trsf = location.Transformation()
|
||
|
||
base_index = len(vertices)
|
||
|
||
# 读取当前 Face 的顶点
|
||
for i in range(1, triangulation.NbNodes() + 1):
|
||
pnt = triangulation.Node(i)
|
||
|
||
# 应用局部变换
|
||
pnt.Transform(trsf)
|
||
|
||
vertices.append([
|
||
pnt.X(),
|
||
pnt.Y(),
|
||
pnt.Z()
|
||
])
|
||
|
||
# 读取当前 Face 的三角形
|
||
for i in range(1, triangulation.NbTriangles() + 1):
|
||
tri = triangulation.Triangle(i)
|
||
|
||
n1, n2, n3 = tri.Get()
|
||
|
||
# OCC 索引从 1 开始，numpy / torch 索引从 0 开始
|
||
faces.append([
|
||
base_index + n1 - 1,
|
||
base_index + n2 - 1,
|
||
base_index + n3 - 1
|
||
])
|
||
|
||
triangle_count += 1
|
||
|
||
face_count += 1
|
||
explorer.Next()
|
||
|
||
vertices_np = np.asarray(vertices, dtype=np.float64)
|
||
faces_np = np.asarray(faces, dtype=np.int64)
|
||
|
||
print("CAD 网格提取完成")
|
||
print(f"CAD 面数量: {face_count}")
|
||
print(f"网格顶点: {vertices_np.shape}")
|
||
print(f"网格三角面: {faces_np.shape}")
|
||
|
||
if vertices_np.shape[0] == 0 or faces_np.shape[0] == 0:
|
||
raise RuntimeError(
|
||
"CAD 三角网格提取失败。可以尝试减小 MESH_DEFLECTION，或者检查 STEP 文件是否正常。"
|
||
)
|
||
|
||
return vertices_np, faces_np
|
||
|
||
|
||
# ============================================================
|
||
# 读取点云
|
||
# ============================================================
|
||
|
||
def read_point_cloud(pointcloud_path):
|
||
if o3d is not None:
|
||
pcd = o3d.io.read_point_cloud(pointcloud_path)
|
||
points_np = np.asarray(pcd.points, dtype=np.float64)
|
||
else:
|
||
pcd = trimesh.load(pointcloud_path)
|
||
points_np = np.asarray(pcd.vertices, dtype=np.float64)
|
||
|
||
if points_np.shape[0] == 0:
|
||
raise RuntimeError(f"点云读取失败或点数为 0: {pointcloud_path}")
|
||
|
||
print("点云读取完成")
|
||
print(f"点云点数: {points_np.shape}")
|
||
|
||
return points_np
|
||
|
||
|
||
def read_mesh_file(mesh_path):
|
||
mesh_path = Path(mesh_path)
|
||
|
||
if o3d is not None:
|
||
mesh = o3d.io.read_triangle_mesh(str(mesh_path))
|
||
if mesh.has_vertices() and mesh.has_triangles():
|
||
vertices_np = np.asarray(mesh.vertices, dtype=np.float64)
|
||
faces_np = np.asarray(mesh.triangles, dtype=np.int64)
|
||
print("Mesh 文件读取完成")
|
||
print(f"Mesh 顶点: {vertices_np.shape}")
|
||
print(f"Mesh 三角面: {faces_np.shape}")
|
||
return vertices_np, faces_np
|
||
|
||
mesh = trimesh.load(str(mesh_path), force="mesh")
|
||
vertices_np = np.asarray(mesh.vertices, dtype=np.float64)
|
||
faces_np = np.asarray(mesh.faces, dtype=np.int64)
|
||
|
||
if vertices_np.shape[0] == 0 or faces_np.shape[0] == 0:
|
||
raise RuntimeError(f"Mesh 文件读取失败或为空: {mesh_path}")
|
||
|
||
print("Mesh 文件读取完成")
|
||
print(f"Mesh 顶点: {vertices_np.shape}")
|
||
print(f"Mesh 三角面: {faces_np.shape}")
|
||
|
||
return vertices_np, faces_np
|
||
|
||
|
||
def read_model_mesh(model_path, mesh_deflection):
|
||
model_path = Path(model_path)
|
||
suffix = model_path.suffix.lower()
|
||
|
||
if suffix in {".step", ".stp"}:
|
||
shape = read_step_file(model_path)
|
||
|
||
print("正在对 CAD 模型进行三角化...")
|
||
mesh = BRepMesh_IncrementalMesh(shape, mesh_deflection)
|
||
mesh.Perform()
|
||
|
||
return occ_shape_to_numpy_mesh(shape)
|
||
|
||
if suffix in {".stl", ".obj", ".ply"}:
|
||
return read_mesh_file(model_path)
|
||
|
||
raise ValueError(f"不支持的模型文件格式: {model_path}")
|
||
|
||
|
||
# ============================================================
|
||
# 点到三角形距离
|
||
# ============================================================
|
||
|
||
def point_to_triangle_distance_squared(points, triangles):
|
||
"""
|
||
points:
|
||
[P, 3]
|
||
triangles:
|
||
[T, 3, 3] or [P, K, 3, 3]
|
||
|
||
返回:
|
||
dist2: [P, T] or [P, K]
|
||
"""
|
||
|
||
p = points[:, None, :]
|
||
|
||
if triangles.dim() == 3:
|
||
a = triangles[None, :, 0, :]
|
||
b = triangles[None, :, 1, :]
|
||
c = triangles[None, :, 2, :]
|
||
elif triangles.dim() == 4:
|
||
a = triangles[:, :, 0, :]
|
||
b = triangles[:, :, 1, :]
|
||
c = triangles[:, :, 2, :]
|
||
else:
|
||
raise ValueError(
|
||
f"triangles must have shape [T, 3, 3] or [P, K, 3, 3], got {tuple(triangles.shape)}"
|
||
)
|
||
|
||
ab = b - a
|
||
ac = c - a
|
||
ap = p - a
|
||
|
||
d1 = (ab * ap).sum(dim=-1)
|
||
d2 = (ac * ap).sum(dim=-1)
|
||
|
||
inf = torch.full_like(d1, torch.inf)
|
||
|
||
# 区域 A
|
||
mask = (d1 <= 0.0) & (d2 <= 0.0)
|
||
dist2 = torch.where(mask, (ap * ap).sum(dim=-1), inf)
|
||
|
||
# 区域 B
|
||
bp = p - b
|
||
d3 = (ab * bp).sum(dim=-1)
|
||
d4 = (ac * bp).sum(dim=-1)
|
||
|
||
mask = (d3 >= 0.0) & (d4 <= d3)
|
||
dist2 = torch.minimum(
|
||
dist2,
|
||
torch.where(mask, (bp * bp).sum(dim=-1), inf)
|
||
)
|
||
|
||
# 边 AB
|
||
vc = d1 * d4 - d3 * d2
|
||
mask = (vc <= 0.0) & (d1 >= 0.0) & (d3 <= 0.0)
|
||
|
||
v = d1 / (d1 - d3).clamp_min(torch.finfo(points.dtype).eps)
|
||
closest = a + v[..., None] * ab
|
||
edge_dist2 = ((p - closest) ** 2).sum(dim=-1)
|
||
|
||
dist2 = torch.minimum(
|
||
dist2,
|
||
torch.where(mask, edge_dist2, inf)
|
||
)
|
||
|
||
# 区域 C
|
||
cp = p - c
|
||
d5 = (ab * cp).sum(dim=-1)
|
||
d6 = (ac * cp).sum(dim=-1)
|
||
|
||
mask = (d6 >= 0.0) & (d5 <= d6)
|
||
dist2 = torch.minimum(
|
||
dist2,
|
||
torch.where(mask, (cp * cp).sum(dim=-1), inf)
|
||
)
|
||
|
||
# 边 AC
|
||
vb = d5 * d2 - d1 * d6
|
||
mask = (vb <= 0.0) & (d2 >= 0.0) & (d6 <= 0.0)
|
||
|
||
w = d2 / (d2 - d6).clamp_min(torch.finfo(points.dtype).eps)
|
||
closest = a + w[..., None] * ac
|
||
edge_dist2 = ((p - closest) ** 2).sum(dim=-1)
|
||
|
||
dist2 = torch.minimum(
|
||
dist2,
|
||
torch.where(mask, edge_dist2, inf)
|
||
)
|
||
|
||
# 边 BC
|
||
va = d3 * d6 - d5 * d4
|
||
mask = (va <= 0.0) & ((d4 - d3) >= 0.0) & ((d5 - d6) >= 0.0)
|
||
|
||
w = (d4 - d3) / ((d4 - d3) + (d5 - d6)).clamp_min(
|
||
torch.finfo(points.dtype).eps
|
||
)
|
||
closest = b + w[..., None] * (c - b)
|
||
edge_dist2 = ((p - closest) ** 2).sum(dim=-1)
|
||
|
||
dist2 = torch.minimum(
|
||
dist2,
|
||
torch.where(mask, edge_dist2, inf)
|
||
)
|
||
|
||
# 三角形内部区域
|
||
denom = (va + vb + vc).clamp_min(torch.finfo(points.dtype).eps)
|
||
v = vb / denom
|
||
w = vc / denom
|
||
|
||
closest = a + ab * v[..., None] + ac * w[..., None]
|
||
face_dist2 = ((p - closest) ** 2).sum(dim=-1)
|
||
|
||
return torch.where(torch.isinf(dist2), face_dist2, dist2)
|
||
|
||
|
||
# ============================================================
|
||
# 分块计算最近距离
|
||
# ============================================================
|
||
|
||
def nearest_aabb_candidate_indices(points, triangle_mins, triangle_maxs, candidate_count):
|
||
"""
|
||
参考 test8.py，用点到三角形 AABB 的距离先筛出每个点附近的候选三角形。
|
||
|
||
points:
|
||
[P, 3]
|
||
triangle_mins / triangle_maxs:
|
||
[M, 3]
|
||
|
||
返回:
|
||
candidate_indices: [P, K]
|
||
"""
|
||
|
||
below = (triangle_mins[None, :, :] - points[:, None, :]).clamp_min(0.0)
|
||
above = (points[:, None, :] - triangle_maxs[None, :, :]).clamp_min(0.0)
|
||
aabb_dist2 = ((below + above) ** 2).sum(dim=-1)
|
||
k = min(candidate_count, triangle_mins.shape[0])
|
||
return torch.topk(aabb_dist2, k=k, dim=1, largest=False).indices
|
||
|
||
|
||
def nearest_mesh_distances(
|
||
points,
|
||
triangles,
|
||
point_chunk_size=2048,
|
||
triangle_chunk_size=4096,
|
||
candidate_triangle_count=DEFAULT_CANDIDATE_TRIANGLE_COUNT
|
||
):
|
||
"""
|
||
分块计算点云到三角网格的最近距离。
|
||
|
||
默认使用 test8.py 的 AABB 候选三角形筛选方法：
|
||
1. 先按点到三角形包围盒的距离选出 K 个候选三角形；
|
||
2. 再对候选三角形计算精确点到三角形距离。
|
||
|
||
如果 candidate_triangle_count <= 0，则退回旧的全量三角面分块精确搜索。
|
||
|
||
points:
|
||
[N, 3]
|
||
triangles:
|
||
[M, 3, 3]
|
||
|
||
返回:
|
||
distances: [N]
|
||
"""
|
||
|
||
all_distances = []
|
||
num_points = points.shape[0]
|
||
num_triangles = triangles.shape[0]
|
||
|
||
if candidate_triangle_count > 0:
|
||
triangle_mins = triangles.min(dim=1).values
|
||
triangle_maxs = triangles.max(dim=1).values
|
||
|
||
for p_start in range(0, num_points, point_chunk_size):
|
||
p_end = min(p_start + point_chunk_size, num_points)
|
||
point_chunk = points[p_start:p_end]
|
||
|
||
candidate_idx = nearest_aabb_candidate_indices(
|
||
point_chunk,
|
||
triangle_mins,
|
||
triangle_maxs,
|
||
candidate_triangle_count
|
||
)
|
||
candidate_triangles = triangles[candidate_idx]
|
||
dist2 = point_to_triangle_distance_squared(
|
||
point_chunk,
|
||
candidate_triangles
|
||
)
|
||
all_distances.append(torch.sqrt(dist2.min(dim=1).values))
|
||
|
||
return torch.cat(all_distances, dim=0)
|
||
|
||
for p_start in range(0, num_points, point_chunk_size):
|
||
p_end = min(p_start + point_chunk_size, num_points)
|
||
|
||
point_chunk = points[p_start:p_end]
|
||
|
||
min_dist2 = torch.full(
|
||
(point_chunk.shape[0],),
|
||
torch.inf,
|
||
dtype=points.dtype,
|
||
device=points.device
|
||
)
|
||
|
||
for t_start in range(0, num_triangles, triangle_chunk_size):
|
||
t_end = min(t_start + triangle_chunk_size, num_triangles)
|
||
|
||
triangle_chunk = triangles[t_start:t_end]
|
||
|
||
dist2 = point_to_triangle_distance_squared(
|
||
point_chunk,
|
||
triangle_chunk
|
||
)
|
||
|
||
min_dist2 = torch.minimum(
|
||
min_dist2,
|
||
dist2.min(dim=1).values
|
||
)
|
||
|
||
all_distances.append(torch.sqrt(min_dist2))
|
||
|
||
return torch.cat(all_distances, dim=0)
|
||
|
||
|
||
# ============================================================
|
||
# 热力图显示
|
||
# ============================================================
|
||
|
||
def visualize_point_cloud_heatmap(points_np, distances_np):
|
||
if o3d is None:
|
||
print("未安装 open3d，跳过可视化。")
|
||
return
|
||
|
||
pcd_show = build_heatmap_point_cloud(points_np, distances_np)
|
||
print("色阶规则：深蓝/紫色 = 最近，红色 = 最远")
|
||
|
||
o3d.visualization.draw_geometries(
|
||
[pcd_show],
|
||
window_name="点云到 CAD 距离热力图"
|
||
)
|
||
|
||
|
||
def build_heatmap_point_cloud(points_np, distances_np):
|
||
d_min = distances_np.min()
|
||
d_max = distances_np.max()
|
||
|
||
dist_norm = (distances_np - d_min) / (d_max - d_min + 1e-12)
|
||
|
||
colors = plt.cm.turbo(dist_norm)[:, :3]
|
||
|
||
pcd_show = o3d.geometry.PointCloud()
|
||
pcd_show.points = o3d.utility.Vector3dVector(points_np)
|
||
pcd_show.colors = o3d.utility.Vector3dVector(colors)
|
||
|
||
return pcd_show
|
||
|
||
|
||
def count_task_result_points(task_results):
|
||
total = 0
|
||
for task_id in sorted(task_results, key=task_sort_key):
|
||
result = task_results[task_id]
|
||
for sample_result in result.get("sample_results", []):
|
||
total += int(len(sample_result["points_np"]))
|
||
return total
|
||
|
||
|
||
def read_binary_xyz_ply_header(ply_path):
|
||
with open(ply_path, "rb") as file:
|
||
header_lines = []
|
||
while True:
|
||
line = file.readline()
|
||
if not line:
|
||
raise ValueError(f"Invalid PLY header: {ply_path}")
|
||
header_lines.append(line)
|
||
if line.strip() == b"end_header":
|
||
break
|
||
|
||
header_size = file.tell()
|
||
|
||
text_lines = [
|
||
line.decode("ascii", errors="replace").strip()
|
||
for line in header_lines
|
||
]
|
||
|
||
if len(text_lines) < 3 or text_lines[0] != "ply":
|
||
raise ValueError(f"Not a PLY file: {ply_path}")
|
||
|
||
if text_lines[1] != "format binary_little_endian 1.0":
|
||
raise ValueError(
|
||
"Only binary_little_endian PLY context point clouds can be streamed."
|
||
)
|
||
|
||
vertex_count = None
|
||
vertex_properties = []
|
||
in_vertex = False
|
||
for line in text_lines[2:]:
|
||
parts = line.split()
|
||
if not parts:
|
||
continue
|
||
if parts[0] == "element":
|
||
in_vertex = parts[1] == "vertex"
|
||
if in_vertex:
|
||
vertex_count = int(parts[2])
|
||
continue
|
||
if in_vertex and parts[0] == "property":
|
||
if len(parts) >= 2 and parts[1] == "list":
|
||
raise ValueError("List properties are not supported in streamed PLY.")
|
||
if len(parts) >= 3:
|
||
vertex_properties.append((parts[1], parts[2]))
|
||
|
||
expected = [("float", "x"), ("float", "y"), ("float", "z")]
|
||
if vertex_count is None or vertex_properties[:3] != expected:
|
||
raise ValueError(
|
||
"Context PLY must start with vertex properties: float x, float y, float z."
|
||
)
|
||
if len(vertex_properties) != 3:
|
||
raise ValueError(
|
||
"Context PLY must contain only float x/y/z vertices for streamed merge."
|
||
)
|
||
|
||
return vertex_count, header_size
|
||
|
||
|
||
def write_xyz_rgb_chunk(file, points, colors):
|
||
points = np.asarray(points)
|
||
colors = np.asarray(colors)
|
||
count = len(points)
|
||
if count == 0:
|
||
return
|
||
|
||
vertex_dtype = np.dtype([
|
||
("x", "<f4"),
|
||
("y", "<f4"),
|
||
("z", "<f4"),
|
||
("red", "u1"),
|
||
("green", "u1"),
|
||
("blue", "u1"),
|
||
])
|
||
vertices = np.empty(count, dtype=vertex_dtype)
|
||
vertices["x"] = points[:, 0].astype(np.float32, copy=False)
|
||
vertices["y"] = points[:, 1].astype(np.float32, copy=False)
|
||
vertices["z"] = points[:, 2].astype(np.float32, copy=False)
|
||
|
||
if colors.ndim == 1:
|
||
vertices["red"] = int(colors[0])
|
||
vertices["green"] = int(colors[1])
|
||
vertices["blue"] = int(colors[2])
|
||
else:
|
||
colors = colors.astype(np.uint8, copy=False)
|
||
vertices["red"] = colors[:, 0]
|
||
vertices["green"] = colors[:, 1]
|
||
vertices["blue"] = colors[:, 2]
|
||
|
||
file.write(vertices.tobytes())
|
||
|
||
|
||
def write_context_xyz_ply_stream(input_path, output_file, header_size, point_count):
|
||
record_size = 12
|
||
chunk_points = 1_000_000
|
||
gray = np.array([158, 158, 158], dtype=np.uint8)
|
||
|
||
with open(input_path, "rb") as input_file:
|
||
input_file.seek(header_size)
|
||
remaining = point_count
|
||
while remaining > 0:
|
||
current = min(chunk_points, remaining)
|
||
data = input_file.read(current * record_size)
|
||
if len(data) != current * record_size:
|
||
raise IOError(f"Unexpected end of context PLY: {input_path}")
|
||
points = np.frombuffer(data, dtype="<f4").reshape(-1, 3)
|
||
write_xyz_rgb_chunk(output_file, points, gray)
|
||
remaining -= current
|
||
|
||
|
||
def write_task_results_stream(task_results, output_file):
|
||
chunk_points = 1_000_000
|
||
for task_id in sorted(task_results, key=task_sort_key):
|
||
result = task_results[task_id]
|
||
for sample_result in result.get("sample_results", []):
|
||
points_np = np.asarray(sample_result["points_np"])
|
||
distances_np = np.asarray(sample_result["distances_np"])
|
||
if len(points_np) == 0:
|
||
continue
|
||
|
||
d_min = float(distances_np.min())
|
||
d_max = float(distances_np.max())
|
||
denom = d_max - d_min + 1e-12
|
||
|
||
for start in range(0, len(points_np), chunk_points):
|
||
end = min(start + chunk_points, len(points_np))
|
||
dist_norm = (distances_np[start:end] - d_min) / denom
|
||
colors = (plt.cm.turbo(dist_norm)[:, :3] * 255.0 + 0.5).astype(np.uint8)
|
||
write_xyz_rgb_chunk(output_file, points_np[start:end], colors)
|
||
|
||
|
||
def save_combined_successful_results_stream(
|
||
task_results,
|
||
output_path,
|
||
context_pointcloud_path=None
|
||
):
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
context_point_count = 0
|
||
context_header_size = None
|
||
context_path = Path(context_pointcloud_path) if context_pointcloud_path else None
|
||
if context_path:
|
||
if context_path.exists():
|
||
context_point_count, context_header_size = read_binary_xyz_ply_header(context_path)
|
||
print(f"完整点云将以流式方式加入合并结果: {context_path}")
|
||
else:
|
||
print(f"完整点云不存在，跳过: {context_path}")
|
||
context_path = None
|
||
|
||
task_point_count = count_task_result_points(task_results)
|
||
total_point_count = context_point_count + task_point_count
|
||
if total_point_count == 0:
|
||
print("没有成功执行的可视化结果，跳过合并点云保存。")
|
||
return None
|
||
|
||
header = (
|
||
"ply\n"
|
||
"format binary_little_endian 1.0\n"
|
||
f"element vertex {total_point_count}\n"
|
||
"property float x\n"
|
||
"property float y\n"
|
||
"property float z\n"
|
||
"property uchar red\n"
|
||
"property uchar green\n"
|
||
"property uchar blue\n"
|
||
"end_header\n"
|
||
)
|
||
|
||
with open(output_path, "wb") as output_file:
|
||
output_file.write(header.encode("ascii"))
|
||
if context_path:
|
||
write_context_xyz_ply_stream(
|
||
context_path,
|
||
output_file,
|
||
context_header_size,
|
||
context_point_count
|
||
)
|
||
write_task_results_stream(task_results, output_file)
|
||
|
||
print(f"合并点云已保存: {output_path}")
|
||
return output_path
|
||
|
||
|
||
def build_combined_successful_point_cloud(
|
||
task_results,
|
||
context_pointcloud_path="datasets/train/train_2_10M.ply",
|
||
arrange_grid=False
|
||
):
|
||
if o3d is None:
|
||
return None
|
||
|
||
task_geometries = []
|
||
for task_id in sorted(task_results, key=task_sort_key):
|
||
result = task_results[task_id]
|
||
for sample_result in result.get("sample_results", []):
|
||
pcd_show = build_heatmap_point_cloud(
|
||
sample_result["points_np"],
|
||
sample_result["distances_np"]
|
||
)
|
||
task_geometries.append(pcd_show)
|
||
|
||
if not task_geometries:
|
||
return None
|
||
|
||
geometries = []
|
||
if context_pointcloud_path:
|
||
context_path = Path(context_pointcloud_path)
|
||
if context_path.exists():
|
||
context = o3d.io.read_point_cloud(str(context_path))
|
||
if len(context.points) > 0:
|
||
if not context.has_colors():
|
||
context.paint_uniform_color([0.62, 0.62, 0.62])
|
||
geometries.append(context)
|
||
print(f"完整点云已加入合并结果: {context_path}")
|
||
else:
|
||
print(f"完整点云不存在，跳过: {context_path}")
|
||
|
||
geometries.extend(task_geometries)
|
||
|
||
if arrange_grid and len(geometries) > 1:
|
||
extents = [
|
||
np.asarray(geometry.get_max_bound()) - np.asarray(geometry.get_min_bound())
|
||
for geometry in geometries
|
||
]
|
||
max_extent = max(float(np.max(extent)) for extent in extents)
|
||
spacing = max(max_extent * 1.35, 1.0)
|
||
cols = int(np.ceil(np.sqrt(len(geometries))))
|
||
|
||
for index, geometry in enumerate(geometries):
|
||
row, col = divmod(index, cols)
|
||
center = (
|
||
np.asarray(geometry.get_min_bound())
|
||
+ np.asarray(geometry.get_max_bound())
|
||
) / 2.0
|
||
target_center = np.array([col * spacing, -row * spacing, 0.0])
|
||
geometry.translate(target_center - center)
|
||
|
||
combined = o3d.geometry.PointCloud()
|
||
for geometry in geometries:
|
||
if not geometry.has_colors():
|
||
geometry.paint_uniform_color([0.62, 0.62, 0.62])
|
||
combined += geometry
|
||
return combined
|
||
|
||
|
||
def save_all_successful_results(
|
||
task_results,
|
||
output_path="vis/combined_contour.ply",
|
||
context_pointcloud_path="datasets/train/train_2_10M.ply"
|
||
):
|
||
try:
|
||
return save_combined_successful_results_stream(
|
||
task_results,
|
||
output_path=output_path,
|
||
context_pointcloud_path=context_pointcloud_path
|
||
)
|
||
except Exception as exc:
|
||
context_path = Path(context_pointcloud_path) if context_pointcloud_path else None
|
||
if context_path and context_path.exists() and context_path.stat().st_size > 1024 ** 3:
|
||
raise RuntimeError(
|
||
f"大文件流式合并失败，无法回退到 Open3D 整体加载: {exc}"
|
||
) from exc
|
||
print(f"流式合并失败，回退到 Open3D 合并: {exc}")
|
||
|
||
if o3d is None:
|
||
print("未安装 open3d，跳过合并点云保存。")
|
||
return None
|
||
|
||
combined = build_combined_successful_point_cloud(
|
||
task_results,
|
||
context_pointcloud_path=context_pointcloud_path
|
||
)
|
||
if combined is None or len(combined.points) == 0:
|
||
print("没有成功执行的可视化结果，跳过合并点云保存。")
|
||
return None
|
||
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
ok = o3d.io.write_point_cloud(str(output_path), combined)
|
||
if not ok:
|
||
raise IOError(f"合并点云保存失败: {output_path}")
|
||
|
||
print(f"合并点云已保存: {output_path}")
|
||
return output_path
|
||
|
||
|
||
def visualize_all_successful_results(
|
||
task_results,
|
||
context_pointcloud_path="datasets/train/train_2_10M.ply"
|
||
):
|
||
if o3d is None:
|
||
print("未安装 open3d，跳过可视化。")
|
||
return
|
||
|
||
combined = build_combined_successful_point_cloud(
|
||
task_results,
|
||
context_pointcloud_path=context_pointcloud_path
|
||
)
|
||
if combined is None or len(combined.points) == 0:
|
||
print("没有成功执行的可视化结果。")
|
||
return
|
||
|
||
print("色阶规则：深蓝/紫色 = 最近，红色 = 最远")
|
||
o3d.visualization.draw_geometries(
|
||
[combined],
|
||
window_name="批量点云到 CAD 距离热力图",
|
||
width=1280,
|
||
height=900
|
||
)
|
||
|
||
|
||
def visualize_task_result(task_results, task_id, sample_id=None):
|
||
"""
|
||
task_results: run_batch_tasks 返回的字典，key 为任务号。
|
||
task_id: 需要可视化的任务号。
|
||
sample_id: 可选样本名，例如 13_2。不传时显示该任务第一个成功样本。
|
||
"""
|
||
|
||
result = task_results.get(str(task_id)) or task_results.get(int(task_id))
|
||
|
||
if result is None:
|
||
raise KeyError(f"未找到任务 {task_id} 的计算结果。")
|
||
|
||
if "sample_results" in result:
|
||
sample_results = result["sample_results"]
|
||
if not sample_results:
|
||
raise KeyError(f"任务 {task_id} 没有可视化数据。")
|
||
|
||
if sample_id is None:
|
||
sample_result = sample_results[0]
|
||
else:
|
||
sample_result = next(
|
||
(
|
||
item for item in sample_results
|
||
if item.get("sample_id") == str(sample_id)
|
||
),
|
||
None
|
||
)
|
||
if sample_result is None:
|
||
raise KeyError(f"任务 {task_id} 中未找到样本 {sample_id}。")
|
||
|
||
visualize_point_cloud_heatmap(
|
||
sample_result["points_np"],
|
||
sample_result["distances_np"]
|
||
)
|
||
return
|
||
|
||
visualize_point_cloud_heatmap(
|
||
result["points_np"],
|
||
result["distances_np"]
|
||
)
|
||
|
||
|
||
def _xlsx_cell_ref(row_index, col_index):
|
||
letters = ""
|
||
col = col_index
|
||
|
||
while col:
|
||
col, remainder = divmod(col - 1, 26)
|
||
letters = chr(65 + remainder) + letters
|
||
|
||
return f"{letters}{row_index}"
|
||
|
||
|
||
def _xlsx_cell_xml(row_index, col_index, value):
|
||
cell_ref = _xlsx_cell_ref(row_index, col_index)
|
||
|
||
if isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool):
|
||
return f'<c r="{cell_ref}"><v>{value}</v></c>'
|
||
|
||
text = escape("" if value is None else str(value))
|
||
return f'<c r="{cell_ref}" t="inlineStr"><is><t>{text}</t></is></c>'
|
||
|
||
|
||
def _xlsx_col_to_index(cell_ref):
|
||
letters = "".join(ch for ch in cell_ref if ch.isalpha())
|
||
index = 0
|
||
for letter in letters:
|
||
index = index * 26 + ord(letter.upper()) - ord("A") + 1
|
||
return index
|
||
|
||
|
||
def _xlsx_rel_target_to_path(target):
|
||
target = target.replace("\\", "/").lstrip("/")
|
||
if target.startswith("xl/"):
|
||
return target
|
||
return posixpath.normpath(posixpath.join("xl", target))
|
||
|
||
|
||
def _xlsx_cell_value(cell, shared_strings, ns):
|
||
value_node = cell.find("a:v", ns)
|
||
inline_node = cell.find("a:is/a:t", ns)
|
||
cell_type = cell.attrib.get("t")
|
||
|
||
if cell_type == "s" and value_node is not None:
|
||
return shared_strings[int(value_node.text)]
|
||
if cell_type == "inlineStr" and inline_node is not None:
|
||
return inline_node.text or ""
|
||
if value_node is not None:
|
||
return value_node.text or ""
|
||
return ""
|
||
|
||
|
||
def read_measurement_template(template_path):
|
||
template_path = Path(template_path)
|
||
ns = {
|
||
"a": "http://schemas.openxmlformats.org/spreadsheetml/2006/main",
|
||
"r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
|
||
}
|
||
|
||
with ZipFile(template_path) as xlsx:
|
||
names = xlsx.namelist()
|
||
|
||
shared_strings = []
|
||
if "xl/sharedStrings.xml" in names:
|
||
shared_root = ET.fromstring(xlsx.read("xl/sharedStrings.xml"))
|
||
for item in shared_root.findall("a:si", ns):
|
||
shared_strings.append(
|
||
"".join((text.text or "") for text in item.findall(".//a:t", ns))
|
||
)
|
||
|
||
workbook_root = ET.fromstring(xlsx.read("xl/workbook.xml"))
|
||
first_sheet = workbook_root.find("a:sheets/a:sheet", ns)
|
||
if first_sheet is None:
|
||
raise RuntimeError(f"模板中没有工作表: {template_path}")
|
||
|
||
sheet_name = first_sheet.attrib.get("name", "Sheet1")
|
||
relation_id = first_sheet.attrib[
|
||
"{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"
|
||
]
|
||
|
||
rels_root = ET.fromstring(xlsx.read("xl/_rels/workbook.xml.rels"))
|
||
relation_targets = {
|
||
rel.attrib["Id"]: rel.attrib["Target"]
|
||
for rel in rels_root
|
||
}
|
||
sheet_path = _xlsx_rel_target_to_path(relation_targets[relation_id])
|
||
|
||
sheet_root = ET.fromstring(xlsx.read(sheet_path))
|
||
parsed_rows = []
|
||
max_col_count = 0
|
||
|
||
for row in sheet_root.findall("a:sheetData/a:row", ns):
|
||
values_by_col = {}
|
||
for cell in row.findall("a:c", ns):
|
||
col_index = _xlsx_col_to_index(cell.attrib["r"])
|
||
values_by_col[col_index] = _xlsx_cell_value(cell, shared_strings, ns)
|
||
max_col_count = max(max_col_count, col_index)
|
||
|
||
if values_by_col:
|
||
parsed_rows.append(values_by_col)
|
||
|
||
table = []
|
||
for values_by_col in parsed_rows:
|
||
table.append([
|
||
values_by_col.get(col_index, "")
|
||
for col_index in range(1, max_col_count + 1)
|
||
])
|
||
|
||
if not table:
|
||
raise RuntimeError(f"模板为空: {template_path}")
|
||
|
||
return sheet_name, table[0], table[1:]
|
||
|
||
|
||
def _normalize_template_task_id(value):
|
||
text = str(value).strip()
|
||
if text.endswith(".0") and text[:-2].isdigit():
|
||
return text[:-2]
|
||
return text
|
||
|
||
|
||
def _find_header_index(headers, candidates):
|
||
normalized = [str(header).strip() for header in headers]
|
||
for candidate in candidates:
|
||
if candidate in normalized:
|
||
return normalized.index(candidate)
|
||
raise ValueError(f"模板缺少列: {'/'.join(candidates)}")
|
||
|
||
|
||
def _row_measurement_conclusion(row):
|
||
mean_dist = row.get("mean_dist")
|
||
if mean_dist is None:
|
||
return "失败"
|
||
try:
|
||
if not np.isfinite(float(mean_dist)):
|
||
return "失败"
|
||
except (TypeError, ValueError):
|
||
return "失败"
|
||
return "成功"
|
||
|
||
|
||
def _build_measurement_table_from_template(rows, template_path):
|
||
template_sheet_name, headers, template_rows = read_measurement_template(template_path)
|
||
task_col = _find_header_index(headers, ["任务号", "任务编号", "task_id"])
|
||
result_col = _find_header_index(headers, ["测量结果"])
|
||
conclusion_col = _find_header_index(headers, ["测量结论", "结论", "状态"])
|
||
|
||
result_by_task = {
|
||
_normalize_template_task_id(row.get("task_id")): row
|
||
for row in rows
|
||
}
|
||
used_task_ids = set()
|
||
table = [headers]
|
||
|
||
for template_row in template_rows:
|
||
output_row = list(template_row)
|
||
if len(output_row) < len(headers):
|
||
output_row.extend([""] * (len(headers) - len(output_row)))
|
||
|
||
task_id = _normalize_template_task_id(output_row[task_col])
|
||
result_row = result_by_task.get(task_id)
|
||
if result_row is None:
|
||
continue
|
||
|
||
output_row[result_col] = result_row.get("mean_dist")
|
||
output_row[conclusion_col] = _row_measurement_conclusion(result_row)
|
||
table.append(output_row)
|
||
used_task_ids.add(task_id)
|
||
|
||
for task_id, result_row in result_by_task.items():
|
||
if task_id in used_task_ids:
|
||
continue
|
||
|
||
output_row = [""] * len(headers)
|
||
output_row[task_col] = task_id
|
||
output_row[result_col] = result_row.get("mean_dist")
|
||
output_row[conclusion_col] = _row_measurement_conclusion(result_row)
|
||
table.append(output_row)
|
||
|
||
return template_sheet_name, table
|
||
|
||
|
||
def write_flatness_excel(
|
||
rows,
|
||
output_path,
|
||
sheet_name="平整度",
|
||
result_template_path=DEFAULT_RESULT_TEMPLATE
|
||
):
|
||
"""
|
||
使用 Python 标准库生成 .xlsx，避免额外依赖。
|
||
"""
|
||
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
if result_template_path and Path(result_template_path).exists():
|
||
sheet_name, table = _build_measurement_table_from_template(
|
||
rows,
|
||
result_template_path
|
||
)
|
||
else:
|
||
headers = ["任务号", "样本数", "平均距离(mm)", "最大距离(mm)", "状态", "说明"]
|
||
table = [headers]
|
||
|
||
for row in rows:
|
||
table.append([
|
||
row.get("task_id"),
|
||
row.get("sample_count"),
|
||
row.get("mean_dist"),
|
||
row.get("max_dist"),
|
||
row.get("status", ""),
|
||
row.get("message", "")
|
||
])
|
||
|
||
sheet_rows = []
|
||
for row_index, row_values in enumerate(table, start=1):
|
||
cells = [
|
||
_xlsx_cell_xml(row_index, col_index, value)
|
||
for col_index, value in enumerate(row_values, start=1)
|
||
]
|
||
sheet_rows.append(f'<row r="{row_index}">{"".join(cells)}</row>')
|
||
|
||
safe_sheet_name = escape(sheet_name[:31])
|
||
|
||
workbook_xml = (
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
||
'<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
|
||
'xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">'
|
||
f'<sheets><sheet name="{safe_sheet_name}" sheetId="1" r:id="rId1"/></sheets>'
|
||
'</workbook>'
|
||
)
|
||
|
||
worksheet_xml = (
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
||
'<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
||
'<sheetData>'
|
||
f'{"".join(sheet_rows)}'
|
||
'</sheetData>'
|
||
'</worksheet>'
|
||
)
|
||
|
||
content_types_xml = (
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
||
'<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
|
||
'<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
|
||
'<Default Extension="xml" ContentType="application/xml"/>'
|
||
'<Override PartName="/xl/workbook.xml" '
|
||
'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>'
|
||
'<Override PartName="/xl/worksheets/sheet1.xml" '
|
||
'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
|
||
'</Types>'
|
||
)
|
||
|
||
root_rels_xml = (
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
||
'<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
|
||
'<Relationship Id="rId1" '
|
||
'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" '
|
||
'Target="xl/workbook.xml"/>'
|
||
'</Relationships>'
|
||
)
|
||
|
||
workbook_rels_xml = (
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
||
'<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
|
||
'<Relationship Id="rId1" '
|
||
'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" '
|
||
'Target="worksheets/sheet1.xml"/>'
|
||
'</Relationships>'
|
||
)
|
||
|
||
with ZipFile(output_path, "w", ZIP_DEFLATED) as xlsx:
|
||
xlsx.writestr("[Content_Types].xml", content_types_xml)
|
||
xlsx.writestr("_rels/.rels", root_rels_xml)
|
||
xlsx.writestr("xl/workbook.xml", workbook_xml)
|
||
xlsx.writestr("xl/_rels/workbook.xml.rels", workbook_rels_xml)
|
||
xlsx.writestr("xl/worksheets/sheet1.xml", worksheet_xml)
|
||
|
||
return output_path
|
||
|
||
|
||
# ============================================================
|
||
# 主函数
|
||
# ============================================================
|
||
|
||
def compute_point_CAD_distance(
|
||
CAD_path,
|
||
pointcloud_path,
|
||
mesh_deflection=0.1,
|
||
point_chunk_size=2048,
|
||
triangle_chunk_size=4096,
|
||
candidate_triangle_count=DEFAULT_CANDIDATE_TRIANGLE_COUNT,
|
||
visualize=False
|
||
):
|
||
print("计算设备:", device)
|
||
print("数据精度:", dtype)
|
||
|
||
# 读取 STEP CAD 或 STL/OBJ/PLY Mesh，并统一转换为 numpy 三角网格。
|
||
vertices_np, faces_np = read_model_mesh(CAD_path, mesh_deflection)
|
||
|
||
# 读取点云
|
||
points_np = read_point_cloud(pointcloud_path)
|
||
|
||
# numpy -> torch
|
||
points = torch.as_tensor(points_np, dtype=dtype, device=device)
|
||
vertices = torch.as_tensor(vertices_np, dtype=dtype, device=device)
|
||
faces = torch.as_tensor(faces_np, dtype=torch.long, device=device)
|
||
|
||
triangles = vertices[faces]
|
||
|
||
print("开始计算点云到模型表面的最近距离...")
|
||
print("点云张量:", points.shape)
|
||
print("三角面张量:", triangles.shape)
|
||
print("候选三角面数量:", candidate_triangle_count)
|
||
|
||
with torch.no_grad():
|
||
distances_torch = nearest_mesh_distances(
|
||
points,
|
||
triangles,
|
||
point_chunk_size=point_chunk_size,
|
||
triangle_chunk_size=triangle_chunk_size,
|
||
candidate_triangle_count=candidate_triangle_count
|
||
)
|
||
|
||
distances_np = distances_torch.detach().cpu().numpy()
|
||
|
||
mean_dist = float(distances_np.mean())
|
||
max_dist = float(distances_np.max())
|
||
min_dist = float(distances_np.min())
|
||
|
||
# print("Distance computation finished")
|
||
print(f"最小距离: {min_dist:.8f}")
|
||
print(f"平均距离: {mean_dist:.8f}")
|
||
print(f"最大距离: {max_dist:.8f}")
|
||
|
||
if visualize:
|
||
visualize_point_cloud_heatmap(points_np, distances_np)
|
||
|
||
return {
|
||
"mean_dist": mean_dist,
|
||
"max_dist": max_dist,
|
||
"min_dist": min_dist,
|
||
"points_np": points_np,
|
||
"distances_np": distances_np
|
||
}
|
||
|
||
|
||
def task_id_from_sample_name(sample_name):
|
||
return str(sample_name).split("_", 1)[0]
|
||
|
||
|
||
def normalize_task_id(task_id):
|
||
return str(task_id).strip()
|
||
|
||
|
||
def task_sort_key(value):
|
||
base = task_id_from_sample_name(value)
|
||
if str(base).isdigit():
|
||
return (0, int(base), str(value))
|
||
return (1, str(base), str(value))
|
||
|
||
|
||
def pointcloud_path_for_sample(data_root, sample_id, pointcloud_kind):
|
||
data_root = Path(data_root)
|
||
|
||
if pointcloud_kind == "clean":
|
||
return data_root / "PointCloud" / f"{sample_id}.ply"
|
||
if pointcloud_kind == "noise":
|
||
return data_root / "Noise_PointCloud" / f"{sample_id}_noisy.ply"
|
||
|
||
raise ValueError("pointcloud_kind must be 'noise' or 'clean'.")
|
||
|
||
|
||
def mesh_sample_paths_for_task(data_root, requested_task):
|
||
data_root = Path(data_root)
|
||
mesh_dir = data_root / "Mesh"
|
||
requested_task = normalize_task_id(requested_task)
|
||
|
||
if not mesh_dir.exists():
|
||
return []
|
||
|
||
mesh_paths = []
|
||
for pattern in ("*.stl", "*.obj", "*.ply"):
|
||
mesh_paths.extend(mesh_dir.glob(pattern))
|
||
|
||
if "_" in requested_task:
|
||
matched = [
|
||
path for path in mesh_paths
|
||
if path.stem == requested_task
|
||
]
|
||
else:
|
||
matched = [
|
||
path for path in mesh_paths
|
||
if task_id_from_sample_name(path.stem) == requested_task
|
||
]
|
||
|
||
return sorted(
|
||
matched,
|
||
key=lambda path: task_sort_key(path.stem)
|
||
)
|
||
|
||
|
||
def get_task_paths(task_id, data_root="High_Speed_Data", pointcloud_kind="noise"):
|
||
data_root = Path(data_root)
|
||
task_id = normalize_task_id(task_id)
|
||
|
||
mesh_paths = mesh_sample_paths_for_task(data_root, task_id)
|
||
if mesh_paths:
|
||
return [
|
||
{
|
||
"task_id": task_id_from_sample_name(mesh_path.stem),
|
||
"sample_id": mesh_path.stem,
|
||
"model_path": mesh_path,
|
||
"pointcloud_path": pointcloud_path_for_sample(
|
||
data_root,
|
||
mesh_path.stem,
|
||
pointcloud_kind
|
||
),
|
||
}
|
||
for mesh_path in mesh_paths
|
||
]
|
||
|
||
sample_id = task_id
|
||
return [
|
||
{
|
||
"task_id": task_id_from_sample_name(sample_id),
|
||
"sample_id": sample_id,
|
||
"model_path": data_root / "CAD" / f"{sample_id}.step",
|
||
"pointcloud_path": pointcloud_path_for_sample(
|
||
data_root,
|
||
sample_id,
|
||
pointcloud_kind
|
||
),
|
||
}
|
||
]
|
||
|
||
def run_batch_tasks(
|
||
task_ids=DEFAULT_TASK_IDS,
|
||
data_root="High_Speed_Data",
|
||
pointcloud_kind="noise",
|
||
excel_output="平整度.xlsx",
|
||
mesh_deflection=0.1,
|
||
point_chunk_size=2048,
|
||
triangle_chunk_size=4096,
|
||
candidate_triangle_count=DEFAULT_CANDIDATE_TRIANGLE_COUNT,
|
||
result_template_path=DEFAULT_RESULT_TEMPLATE
|
||
):
|
||
global task_distance_data
|
||
|
||
task_summaries = {}
|
||
task_distance_data = {}
|
||
|
||
for task_id in task_ids:
|
||
task_inputs = get_task_paths(
|
||
task_id,
|
||
data_root=data_root,
|
||
pointcloud_kind=pointcloud_kind
|
||
)
|
||
|
||
for task_input in task_inputs:
|
||
base_task_id = task_input["task_id"]
|
||
sample_id = task_input["sample_id"]
|
||
model_path = task_input["model_path"]
|
||
pointcloud_path = task_input["pointcloud_path"]
|
||
|
||
summary = task_summaries.setdefault(
|
||
base_task_id,
|
||
{
|
||
"task_id": base_task_id,
|
||
"sample_ids": [],
|
||
"point_count": 0,
|
||
"distance_sum": 0.0,
|
||
"max_dist": None,
|
||
"failures": [],
|
||
"sample_results": [],
|
||
}
|
||
)
|
||
summary["sample_ids"].append(sample_id)
|
||
|
||
print(f"===== 任务 {base_task_id} / 样本 {sample_id} =====")
|
||
|
||
missing_paths = [
|
||
str(path)
|
||
for path in [model_path, pointcloud_path]
|
||
if not path.exists()
|
||
]
|
||
|
||
if missing_paths:
|
||
message = f"{sample_id} 缺少文件: " + "; ".join(missing_paths)
|
||
print(message)
|
||
summary["failures"].append(message)
|
||
continue
|
||
|
||
try:
|
||
result = compute_point_CAD_distance(
|
||
CAD_path=str(model_path),
|
||
pointcloud_path=str(pointcloud_path),
|
||
mesh_deflection=mesh_deflection,
|
||
point_chunk_size=point_chunk_size,
|
||
triangle_chunk_size=triangle_chunk_size,
|
||
candidate_triangle_count=candidate_triangle_count,
|
||
visualize=False
|
||
)
|
||
except Exception as exc:
|
||
message = f"{sample_id} 失败: {exc}"
|
||
print(message)
|
||
summary["failures"].append(message)
|
||
continue
|
||
|
||
distances_np = result["distances_np"]
|
||
point_count = len(distances_np)
|
||
summary["point_count"] += point_count
|
||
summary["distance_sum"] += float(distances_np.sum())
|
||
summary["max_dist"] = (
|
||
result["max_dist"]
|
||
if summary["max_dist"] is None
|
||
else max(summary["max_dist"], result["max_dist"])
|
||
)
|
||
summary["sample_results"].append({
|
||
"sample_id": sample_id,
|
||
"model_path": str(model_path),
|
||
"pointcloud_path": str(pointcloud_path),
|
||
"points_np": result["points_np"],
|
||
"distances_np": distances_np,
|
||
"mean_dist": result["mean_dist"],
|
||
"max_dist": result["max_dist"],
|
||
"min_dist": result["min_dist"],
|
||
})
|
||
|
||
result_rows = []
|
||
for task_id in sorted(
|
||
task_summaries,
|
||
key=task_sort_key
|
||
):
|
||
summary = task_summaries[task_id]
|
||
success_count = len(summary["sample_results"])
|
||
failure_count = len(summary["failures"])
|
||
|
||
if success_count:
|
||
mean_dist = summary["distance_sum"] / summary["point_count"]
|
||
max_dist = summary["max_dist"]
|
||
status = "完成" if failure_count == 0 else "部分完成"
|
||
else:
|
||
mean_dist = None
|
||
max_dist = None
|
||
status = "失败"
|
||
|
||
message_parts = []
|
||
if summary["sample_ids"]:
|
||
message_parts.append("样本: " + ",".join(summary["sample_ids"]))
|
||
if summary["failures"]:
|
||
message_parts.append("失败: " + " | ".join(summary["failures"]))
|
||
|
||
result_rows.append({
|
||
"task_id": task_id,
|
||
"sample_count": success_count,
|
||
"mean_dist": mean_dist,
|
||
"max_dist": max_dist,
|
||
"status": status,
|
||
"message": "; ".join(message_parts)
|
||
})
|
||
|
||
task_distance_data[task_id] = {
|
||
"sample_results": summary["sample_results"],
|
||
"sample_ids": [item["sample_id"] for item in summary["sample_results"]],
|
||
}
|
||
|
||
excel_path = write_flatness_excel(
|
||
result_rows,
|
||
excel_output,
|
||
sheet_name="平整度",
|
||
result_template_path=result_template_path
|
||
)
|
||
|
||
print(f"Excel 已保存: {excel_path}")
|
||
|
||
return result_rows, task_distance_data
|
||
|
||
|
||
# ============================================================
|
||
# 命令行入口
|
||
# ============================================================
|
||
|
||
if __name__ == "__main__":
|
||
start_time = time.perf_counter()
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="计算点云到 CAD STEP 模型的最近距离，支持单任务或批量任务。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--CAD",
|
||
"-c",
|
||
default="High_Speed_Data/CAD/11.step",
|
||
help="CAD STEP 文件路径，仅在 --single 单任务模式下使用。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--pointcloud",
|
||
"-p",
|
||
default="High_Speed_Data/Noise_PointCloud/11_noisy.ply",
|
||
help="点云 PLY 文件路径，仅在 --single 单任务模式下使用。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--tasks",
|
||
default=",".join(str(task_id) for task_id in DEFAULT_TASK_IDS),
|
||
help="批量任务号或样本名，多个值用英文逗号分隔。例：13 会匹配 13_1/13_2；13_1 只匹配该样本。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--data_root",
|
||
default="High_Speed_Data",
|
||
help="批量任务的数据根目录。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--pointcloud_kind",
|
||
choices=["noise", "clean"],
|
||
default="noise",
|
||
help="批量任务点云来源：noise=Noise_PointCloud，clean=PointCloud。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--excel_output",
|
||
default="平整度.xlsx",
|
||
help="批量结果 Excel 输出路径。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--result_template",
|
||
default=str(DEFAULT_RESULT_TEMPLATE),
|
||
help="测量结果模板 Excel 路径。新表会按任务号复制模板对应行，并写入测量结果和测量结论。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--combined_output",
|
||
default="vis/combined_contour.ply",
|
||
help="批量模式下所有成功热力图合并后的 PLY 输出路径。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--combined_context_pointcloud",
|
||
default="datasets/train/train_2_10M.ply",
|
||
help="批量合并热力图时一起加入的完整点云路径，可切换为 train_2_1T.ply 或 train_2_100M.ply。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--single",
|
||
action="store_true",
|
||
help="使用 --CAD 和 --pointcloud 执行单任务，而不是默认批量任务。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--visualize",
|
||
action="store_true",
|
||
help="显示热力图；批量模式会把所有成功结果合并到一个窗口。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--mesh_deflection",
|
||
default=0.1,
|
||
type=float,
|
||
help="CAD 三角化精度，越小越精细，但计算越慢。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--point_chunk_size",
|
||
default=2048,
|
||
type=int,
|
||
help="点云分块大小，显存不够时可调小。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--triangle_chunk_size",
|
||
default=4096,
|
||
type=int,
|
||
help="全量精确搜索模式下的三角面分块大小；candidate_triangle_count <= 0 时使用。"
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--candidate_triangle_count",
|
||
default=DEFAULT_CANDIDATE_TRIANGLE_COUNT,
|
||
type=int,
|
||
help="每个点按 AABB 预筛的候选三角面数量；设为 0 或负数则使用全量精确搜索。"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.single:
|
||
result = compute_point_CAD_distance(
|
||
CAD_path=args.CAD,
|
||
pointcloud_path=args.pointcloud,
|
||
mesh_deflection=args.mesh_deflection,
|
||
point_chunk_size=args.point_chunk_size,
|
||
triangle_chunk_size=args.triangle_chunk_size,
|
||
candidate_triangle_count=args.candidate_triangle_count,
|
||
visualize=args.visualize
|
||
)
|
||
|
||
print("===== 计算结果 =====")
|
||
print(f"平均距离: {result['mean_dist']:.8f} mm")
|
||
print(f"最大距离: {result['max_dist']:.8f} mm")
|
||
elapsed_time = time.perf_counter() - start_time
|
||
print(f"运行时间: {elapsed_time:.2f} 秒")
|
||
else:
|
||
task_ids = [
|
||
task_id.strip()
|
||
for task_id in args.tasks.split(",")
|
||
if task_id.strip()
|
||
]
|
||
|
||
result_rows, task_distance_data = run_batch_tasks(
|
||
task_ids=task_ids,
|
||
data_root=args.data_root,
|
||
pointcloud_kind=args.pointcloud_kind,
|
||
excel_output=args.excel_output,
|
||
mesh_deflection=args.mesh_deflection,
|
||
point_chunk_size=args.point_chunk_size,
|
||
triangle_chunk_size=args.triangle_chunk_size,
|
||
candidate_triangle_count=args.candidate_triangle_count,
|
||
result_template_path=args.result_template
|
||
)
|
||
|
||
print("===== 批量计算结果 =====")
|
||
for row in result_rows:
|
||
if row["status"] in {"完成", "部分完成"}:
|
||
print(
|
||
f"{row['task_id']}: "
|
||
f"样本数={row['sample_count']}, "
|
||
f"平均距离={row['mean_dist']:.8f} mm, "
|
||
f"最大距离={row['max_dist']:.8f} mm"
|
||
)
|
||
else:
|
||
print(f"{row['task_id']}: {row['status']} - {row['message']}")
|
||
|
||
elapsed_time = time.perf_counter() - start_time
|
||
print(f"运行时间: {elapsed_time:.2f} 秒")
|
||
|
||
print("task_distance_data 已按任务号保存 sample_results。")
|
||
save_all_successful_results(
|
||
task_distance_data,
|
||
output_path=args.combined_output,
|
||
context_pointcloud_path=args.combined_context_pointcloud
|
||
)
|
||
if args.visualize:
|
||
visualize_all_successful_results(
|
||
task_distance_data,
|
||
context_pointcloud_path=args.combined_context_pointcloud
|
||
)
|