447 lines
12 KiB
Python
447 lines
12 KiB
Python
import os
|
|
import argparse
|
|
import numpy as np
|
|
import open3d as o3d
|
|
|
|
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
|
|
|
|
try:
|
|
from OCC.Core.TopoDS import topods
|
|
USE_TOPODS = "topods"
|
|
except ImportError:
|
|
from OCC.Core.TopoDS import TopoDS
|
|
USE_TOPODS = "TopoDS"
|
|
|
|
|
|
# ============================================================
|
|
# STEP -> Open3D Mesh
|
|
# ============================================================
|
|
|
|
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("[OK] STEP loaded:", step_path)
|
|
return shape
|
|
|
|
|
|
def occ_shape_to_numpy_mesh(shape):
|
|
vertices = []
|
|
triangles = []
|
|
|
|
explorer = TopExp_Explorer(shape, TopAbs_FACE)
|
|
|
|
face_count = 0
|
|
tri_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)
|
|
|
|
for i in range(1, triangulation.NbNodes() + 1):
|
|
pnt = triangulation.Node(i)
|
|
pnt.Transform(trsf)
|
|
|
|
vertices.append([
|
|
pnt.X(),
|
|
pnt.Y(),
|
|
pnt.Z()
|
|
])
|
|
|
|
for i in range(1, triangulation.NbTriangles() + 1):
|
|
tri = triangulation.Triangle(i)
|
|
n1, n2, n3 = tri.Get()
|
|
|
|
triangles.append([
|
|
base_index + n1 - 1,
|
|
base_index + n2 - 1,
|
|
base_index + n3 - 1
|
|
])
|
|
|
|
tri_count += 1
|
|
|
|
face_count += 1
|
|
explorer.Next()
|
|
|
|
vertices = np.asarray(vertices, dtype=np.float64)
|
|
triangles = np.asarray(triangles, dtype=np.int32)
|
|
|
|
print("[OK] CAD mesh extracted")
|
|
print(" CAD faces:", face_count)
|
|
print(" mesh vertices:", vertices.shape)
|
|
print(" mesh triangles:", triangles.shape)
|
|
|
|
if vertices.shape[0] == 0 or triangles.shape[0] == 0:
|
|
raise RuntimeError("CAD 三角网格提取失败")
|
|
|
|
return vertices, triangles
|
|
|
|
|
|
def step_to_open3d_mesh(step_path, mesh_deflection=0.1):
|
|
shape = read_step_file(step_path)
|
|
|
|
print("[INFO] Meshing CAD...")
|
|
mesh_algo = BRepMesh_IncrementalMesh(shape, mesh_deflection)
|
|
mesh_algo.Perform()
|
|
|
|
vertices, triangles = occ_shape_to_numpy_mesh(shape)
|
|
|
|
mesh = o3d.geometry.TriangleMesh()
|
|
mesh.vertices = o3d.utility.Vector3dVector(vertices)
|
|
mesh.triangles = o3d.utility.Vector3iVector(triangles)
|
|
|
|
mesh.remove_duplicated_vertices()
|
|
mesh.remove_duplicated_triangles()
|
|
mesh.remove_degenerate_triangles()
|
|
mesh.compute_vertex_normals()
|
|
|
|
return mesh
|
|
|
|
|
|
# ============================================================
|
|
# 点云处理
|
|
# ============================================================
|
|
|
|
def make_translation_matrix(dx, dy, dz):
|
|
T = np.eye(4, dtype=np.float64)
|
|
T[0, 3] = dx
|
|
T[1, 3] = dy
|
|
T[2, 3] = dz
|
|
return T
|
|
|
|
|
|
def crop_front_tail_by_x(pcd, ratio=0.08):
|
|
"""
|
|
按 X 轴裁剪车头和车尾。
|
|
|
|
front: X 最小端
|
|
tail: X 最大端
|
|
"""
|
|
|
|
points = np.asarray(pcd.points)
|
|
|
|
min_bound = points.min(axis=0)
|
|
max_bound = points.max(axis=0)
|
|
|
|
x_min = min_bound[0]
|
|
x_max = max_bound[0]
|
|
x_len = x_max - x_min
|
|
|
|
front_min = min_bound.copy()
|
|
front_max = max_bound.copy()
|
|
front_min[0] = x_min
|
|
front_max[0] = x_min + ratio * x_len
|
|
|
|
tail_min = min_bound.copy()
|
|
tail_max = max_bound.copy()
|
|
tail_min[0] = x_max - ratio * x_len
|
|
tail_max[0] = x_max
|
|
|
|
front_box = o3d.geometry.AxisAlignedBoundingBox(
|
|
min_bound=front_min,
|
|
max_bound=front_max
|
|
)
|
|
|
|
tail_box = o3d.geometry.AxisAlignedBoundingBox(
|
|
min_bound=tail_min,
|
|
max_bound=tail_max
|
|
)
|
|
|
|
front_pcd = pcd.crop(front_box)
|
|
tail_pcd = pcd.crop(tail_box)
|
|
|
|
print("[OK] Crop front/tail by X")
|
|
print(" X range:", x_min, x_max)
|
|
print(" front points:", np.asarray(front_pcd.points).shape)
|
|
print(" tail points:", np.asarray(tail_pcd.points).shape)
|
|
|
|
return front_pcd, tail_pcd
|
|
|
|
|
|
def add_gaussian_noise(pcd, noise_std=0.0):
|
|
"""
|
|
给点云添加高斯噪声。
|
|
noise_std=0 表示不加噪声。
|
|
"""
|
|
|
|
if noise_std <= 0:
|
|
return pcd
|
|
|
|
points = np.asarray(pcd.points)
|
|
noise = np.random.normal(
|
|
loc=0.0,
|
|
scale=noise_std,
|
|
size=points.shape
|
|
)
|
|
|
|
noisy_points = points + noise
|
|
|
|
noisy_pcd = o3d.geometry.PointCloud()
|
|
noisy_pcd.points = o3d.utility.Vector3dVector(noisy_points)
|
|
|
|
if pcd.has_colors():
|
|
noisy_pcd.colors = pcd.colors
|
|
|
|
if pcd.has_normals():
|
|
noisy_pcd.normals = pcd.normals
|
|
|
|
return noisy_pcd
|
|
|
|
|
|
def visualize_pair(cad_pcd, shifted_pcd):
|
|
cad_vis = o3d.geometry.PointCloud(cad_pcd)
|
|
shifted_vis = o3d.geometry.PointCloud(shifted_pcd)
|
|
|
|
cad_vis.paint_uniform_color([0.0, 0.5, 1.0])
|
|
shifted_vis.paint_uniform_color([1.0, 0.0, 0.0])
|
|
|
|
o3d.visualization.draw_geometries(
|
|
[cad_vis, shifted_vis],
|
|
window_name="Blue = CAD sampled point cloud, Red = shifted test point cloud"
|
|
)
|
|
|
|
|
|
# ============================================================
|
|
# 主函数
|
|
# ============================================================
|
|
|
|
def generate_shifted_train_pointcloud(
|
|
cad_path,
|
|
output_dir,
|
|
sample_points=1000000,
|
|
mesh_deflection=0.1,
|
|
crop_ratio=0.08,
|
|
dx=100.0,
|
|
dy=20.0,
|
|
dz=-10.0,
|
|
noise_std=0.0,
|
|
visualize=True
|
|
):
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 1. CAD -> mesh
|
|
cad_mesh = step_to_open3d_mesh(
|
|
cad_path,
|
|
mesh_deflection=mesh_deflection
|
|
)
|
|
|
|
# 2. mesh -> 点云
|
|
print(f"[INFO] Sampling CAD surface: {sample_points} points")
|
|
cad_pcd = cad_mesh.sample_points_uniformly(
|
|
number_of_points=sample_points
|
|
)
|
|
|
|
cad_pcd_path = os.path.join(output_dir, "cad_sampled_whole.ply")
|
|
o3d.io.write_point_cloud(cad_pcd_path, cad_pcd)
|
|
print("[OK] Saved:", cad_pcd_path)
|
|
|
|
# 3. 裁剪原始 CAD 点云的车头 / 车尾
|
|
cad_front, cad_tail = crop_front_tail_by_x(
|
|
cad_pcd,
|
|
ratio=crop_ratio
|
|
)
|
|
|
|
o3d.io.write_point_cloud(
|
|
os.path.join(output_dir, "cad_front.ply"),
|
|
cad_front
|
|
)
|
|
|
|
o3d.io.write_point_cloud(
|
|
os.path.join(output_dir, "cad_tail.ply"),
|
|
cad_tail
|
|
)
|
|
|
|
# 4. 生成固定平移矩阵
|
|
T = make_translation_matrix(dx, dy, dz)
|
|
|
|
print("\n========== Fixed Translation ==========")
|
|
print("dx:", dx)
|
|
print("dy:", dy)
|
|
print("dz:", dz)
|
|
print("T:")
|
|
print(T)
|
|
|
|
# 5. 对整车点云施加固定平移
|
|
shifted_pcd = o3d.geometry.PointCloud(cad_pcd)
|
|
shifted_pcd.transform(T)
|
|
|
|
# 6. 可选：加噪声，模拟真实扫描点云
|
|
shifted_pcd = add_gaussian_noise(
|
|
shifted_pcd,
|
|
noise_std=noise_std
|
|
)
|
|
|
|
shifted_path = os.path.join(output_dir, "shifted_whole_pointcloud.ply")
|
|
o3d.io.write_point_cloud(shifted_path, shifted_pcd)
|
|
print("[OK] Saved:", shifted_path)
|
|
|
|
# 7. 裁剪平移后的车头 / 车尾
|
|
shifted_front, shifted_tail = crop_front_tail_by_x(
|
|
shifted_pcd,
|
|
ratio=crop_ratio
|
|
)
|
|
|
|
shifted_front_path = os.path.join(output_dir, "shifted_front.ply")
|
|
shifted_tail_path = os.path.join(output_dir, "shifted_tail.ply")
|
|
|
|
o3d.io.write_point_cloud(shifted_front_path, shifted_front)
|
|
o3d.io.write_point_cloud(shifted_tail_path, shifted_tail)
|
|
|
|
print("[OK] Saved:", shifted_front_path)
|
|
print("[OK] Saved:", shifted_tail_path)
|
|
|
|
# 8. 保存真实位移矩阵
|
|
matrix_path = os.path.join(output_dir, "GT_translation_matrix.txt")
|
|
np.savetxt(matrix_path, T, fmt="%.10f")
|
|
print("[OK] Saved:", matrix_path)
|
|
|
|
# 9. 保存真实位移向量
|
|
translation_path = os.path.join(output_dir, "GT_translation_vector.txt")
|
|
np.savetxt(
|
|
translation_path,
|
|
np.array([dx, dy, dz], dtype=np.float64),
|
|
fmt="%.10f"
|
|
)
|
|
print("[OK] Saved:", translation_path)
|
|
|
|
# 10. 可视化
|
|
if visualize:
|
|
# 百万级点云显示会卡，所以这里可视化时降采样
|
|
cad_vis = cad_pcd.voxel_down_sample(voxel_size=20.0)
|
|
shifted_vis = shifted_pcd.voxel_down_sample(voxel_size=20.0)
|
|
|
|
visualize_pair(cad_vis, shifted_vis)
|
|
|
|
return {
|
|
"cad_sampled_path": cad_pcd_path,
|
|
"shifted_pointcloud_path": shifted_path,
|
|
"shifted_front_path": shifted_front_path,
|
|
"shifted_tail_path": shifted_tail_path,
|
|
"translation_matrix_path": matrix_path,
|
|
"translation_vector": np.array([dx, dy, dz], dtype=np.float64)
|
|
}
|
|
|
|
|
|
# ============================================================
|
|
# 命令行入口
|
|
# ============================================================
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="从整车 CAD 生成测试点云，并施加固定前后/整体位移"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--cad",
|
|
default=r"datasets/CAD/M02车.stp",
|
|
help="整车 CAD STEP/STP 文件路径"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
default=r"datasets\shift_test_output",
|
|
help="输出目录"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--sample_points",
|
|
type=int,
|
|
default=1000000,
|
|
help="从 CAD 表面采样的点数，默认 1000000"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--mesh_deflection",
|
|
type=float,
|
|
default=0.1,
|
|
help="CAD 三角化精度，越小越精细，默认 0.1"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--crop_ratio",
|
|
type=float,
|
|
default=0.08,
|
|
help="车头/车尾裁剪比例，默认 0.08"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dx",
|
|
type=float,
|
|
default=100.0,
|
|
help="X 方向固定平移量"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dy",
|
|
type=float,
|
|
default=20.0,
|
|
help="Y 方向固定平移量"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dz",
|
|
type=float,
|
|
default=-10.0,
|
|
help="Z 方向固定平移量"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--noise_std",
|
|
type=float,
|
|
default=0.0,
|
|
help="高斯噪声标准差，默认 0 表示不加噪声"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--no_visualize",
|
|
action="store_true",
|
|
help="不显示可视化窗口"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
result = generate_shifted_train_pointcloud(
|
|
cad_path=args.cad,
|
|
output_dir=args.output_dir,
|
|
sample_points=args.sample_points,
|
|
mesh_deflection=args.mesh_deflection,
|
|
crop_ratio=args.crop_ratio,
|
|
dx=args.dx,
|
|
dy=args.dy,
|
|
dz=args.dz,
|
|
noise_std=args.noise_std,
|
|
visualize=not args.no_visualize
|
|
)
|
|
|
|
print("\n========== Result ==========")
|
|
print("CAD sampled point cloud:", result["cad_sampled_path"])
|
|
print("Shifted whole point cloud:", result["shifted_pointcloud_path"])
|
|
print("Shifted front point cloud:", result["shifted_front_path"])
|
|
print("Shifted tail point cloud:", result["shifted_tail_path"])
|
|
print("GT translation vector:", result["translation_vector"])
|
|
print("GT matrix:", result["translation_matrix_path"]) |