800 lines
22 KiB
Python
800 lines
22 KiB
Python
import argparse
|
||
from pathlib import Path
|
||
from zipfile import ZipFile, ZIP_DEFLATED
|
||
from xml.sax.saxutils import escape
|
||
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
|
||
]
|
||
|
||
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 point_to_triangle_distance_squared(points, triangles):
|
||
"""
|
||
points:
|
||
[P, 3]
|
||
triangles:
|
||
[T, 3, 3]
|
||
|
||
返回:
|
||
dist2: [P, T]
|
||
"""
|
||
|
||
p = points[:, None, :]
|
||
|
||
a = triangles[None, :, 0, :]
|
||
b = triangles[None, :, 1, :]
|
||
c = triangles[None, :, 2, :]
|
||
|
||
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)
|
||
|
||
dist2 = torch.minimum(dist2, face_dist2)
|
||
|
||
return dist2
|
||
|
||
|
||
# ============================================================
|
||
# 分块计算最近距离
|
||
# ============================================================
|
||
|
||
def nearest_mesh_distances(
|
||
points,
|
||
triangles,
|
||
point_chunk_size=2048,
|
||
triangle_chunk_size=4096
|
||
):
|
||
"""
|
||
为了避免显存爆炸，点和三角面都分块计算。
|
||
|
||
points:
|
||
[N, 3]
|
||
triangles:
|
||
[M, 3, 3]
|
||
|
||
返回:
|
||
distances: [N]
|
||
"""
|
||
|
||
all_distances = []
|
||
|
||
num_points = points.shape[0]
|
||
num_triangles = triangles.shape[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))
|
||
|
||
# print(f"processed points: {p_end}/{num_points}")
|
||
|
||
distances = torch.cat(all_distances, dim=0)
|
||
|
||
return distances
|
||
|
||
|
||
# ============================================================
|
||
# 热力图显示
|
||
# ============================================================
|
||
|
||
def visualize_point_cloud_heatmap(points_np, distances_np):
|
||
if o3d is None:
|
||
print("未安装 open3d，跳过可视化。")
|
||
return
|
||
|
||
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)
|
||
|
||
print("色阶规则：深蓝/紫色 = 最近，红色 = 最远")
|
||
|
||
o3d.visualization.draw_geometries(
|
||
[pcd_show],
|
||
window_name="点云到 CAD 距离热力图"
|
||
)
|
||
|
||
|
||
def visualize_task_result(task_results, task_id):
|
||
"""
|
||
task_results: run_batch_tasks 返回的字典，key 为任务号。
|
||
task_id: 需要可视化的任务号。
|
||
"""
|
||
|
||
result = task_results.get(int(task_id))
|
||
|
||
if result is None:
|
||
raise KeyError(f"未找到任务 {task_id} 的计算结果。")
|
||
|
||
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 write_flatness_excel(rows, output_path, sheet_name="平整度"):
|
||
"""
|
||
使用 Python 标准库生成 .xlsx，避免额外依赖。
|
||
"""
|
||
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
headers = ["任务号", "平均距离(mm)", "最大距离(mm)", "状态", "说明"]
|
||
table = [headers]
|
||
|
||
for row in rows:
|
||
table.append([
|
||
row.get("task_id"),
|
||
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,
|
||
visualize=False
|
||
):
|
||
print("计算设备:", device)
|
||
print("数据精度:", dtype)
|
||
|
||
# STEP -> OCC Shape
|
||
shape = read_step_file(CAD_path)
|
||
|
||
# OCC Shape 三角化
|
||
print("正在对 CAD 模型进行三角化...")
|
||
mesh = BRepMesh_IncrementalMesh(shape, mesh_deflection)
|
||
mesh.Perform()
|
||
|
||
# 提取 numpy 网格
|
||
vertices_np, faces_np = occ_shape_to_numpy_mesh(shape)
|
||
|
||
# 读取点云
|
||
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("开始计算点云到 CAD 的最近距离...")
|
||
print("点云张量:", points.shape)
|
||
print("三角面张量:", triangles.shape)
|
||
|
||
with torch.no_grad():
|
||
distances_torch = nearest_mesh_distances(
|
||
points,
|
||
triangles,
|
||
point_chunk_size=point_chunk_size,
|
||
triangle_chunk_size=triangle_chunk_size
|
||
)
|
||
|
||
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 get_task_paths(task_id, data_root="High_Speed_Data", pointcloud_kind="noise"):
|
||
data_root = Path(data_root)
|
||
task_id = int(task_id)
|
||
|
||
cad_path = data_root / "CAD" / f"{task_id}.step"
|
||
|
||
if pointcloud_kind == "clean":
|
||
pointcloud_path = data_root / "PointCloud" / f"{task_id}.ply"
|
||
else:
|
||
pointcloud_path = data_root / "Noise_PointCloud" / f"{task_id}_noisy.ply"
|
||
|
||
return cad_path, pointcloud_path
|
||
|
||
|
||
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
|
||
):
|
||
global task_distance_data
|
||
|
||
result_rows = []
|
||
task_distance_data = {}
|
||
|
||
for task_id in task_ids:
|
||
task_id = int(task_id)
|
||
cad_path, pointcloud_path = get_task_paths(
|
||
task_id,
|
||
data_root=data_root,
|
||
pointcloud_kind=pointcloud_kind
|
||
)
|
||
|
||
print(f"===== 任务 {task_id} =====")
|
||
|
||
missing_paths = [
|
||
str(path)
|
||
for path in [cad_path, pointcloud_path]
|
||
if not path.exists()
|
||
]
|
||
|
||
if missing_paths:
|
||
message = "缺少文件: " + "; ".join(missing_paths)
|
||
print(message)
|
||
result_rows.append({
|
||
"task_id": task_id,
|
||
"mean_dist": None,
|
||
"max_dist": None,
|
||
"status": "失败",
|
||
"message": message
|
||
})
|
||
continue
|
||
|
||
try:
|
||
result = compute_point_CAD_distance(
|
||
CAD_path=str(cad_path),
|
||
pointcloud_path=str(pointcloud_path),
|
||
mesh_deflection=mesh_deflection,
|
||
point_chunk_size=point_chunk_size,
|
||
triangle_chunk_size=triangle_chunk_size,
|
||
visualize=False
|
||
)
|
||
except Exception as exc:
|
||
message = str(exc)
|
||
print(f"任务 {task_id} 失败: {message}")
|
||
result_rows.append({
|
||
"task_id": task_id,
|
||
"mean_dist": None,
|
||
"max_dist": None,
|
||
"status": "失败",
|
||
"message": message
|
||
})
|
||
continue
|
||
|
||
task_distance_data[task_id] = {
|
||
"points_np": result["points_np"],
|
||
"distances_np": result["distances_np"]
|
||
}
|
||
|
||
result_rows.append({
|
||
"task_id": task_id,
|
||
"mean_dist": result["mean_dist"],
|
||
"max_dist": result["max_dist"],
|
||
"status": "完成",
|
||
"message": ""
|
||
})
|
||
|
||
excel_path = write_flatness_excel(
|
||
result_rows,
|
||
excel_output,
|
||
sheet_name="平整度"
|
||
)
|
||
|
||
print(f"Excel 已保存: {excel_path}")
|
||
|
||
return result_rows, task_distance_data
|
||
|
||
|
||
# ============================================================
|
||
# 命令行入口
|
||
# ============================================================
|
||
|
||
if __name__ == "__main__":
|
||
|
||
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="批量任务号，多个任务用英文逗号分隔。"
|
||
)
|
||
|
||
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(
|
||
"--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="三角面分块大小，显存不够时可调小。"
|
||
)
|
||
|
||
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,
|
||
visualize=args.visualize
|
||
)
|
||
|
||
print("===== 计算结果 =====")
|
||
print(f"平均距离: {result['mean_dist']:.8f} mm")
|
||
print(f"最大距离: {result['max_dist']:.8f} mm")
|
||
else:
|
||
task_ids = [
|
||
int(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
|
||
)
|
||
|
||
print("===== 批量计算结果 =====")
|
||
for row in result_rows:
|
||
if row["status"] == "完成":
|
||
print(
|
||
f"{row['task_id']}: "
|
||
f"平均距离={row['mean_dist']:.8f} mm, "
|
||
f"最大距离={row['max_dist']:.8f} mm"
|
||
)
|
||
else:
|
||
print(f"{row['task_id']}: {row['status']} - {row['message']}")
|
||
|
||
print("task_distance_data 已按任务号保存 points_np 和 distances_np。")
|