You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
555 lines
26 KiB
555 lines
26 KiB
"""
|
|
CAD模型处理脚本
|
|
功能：将STEP格式的CAD模型转换为结构化数据，包括：
|
|
- 几何信息：面、边、顶点的坐标数据
|
|
- 拓扑信息：面-边-顶点的邻接关系
|
|
- 空间信息：包围盒数据
|
|
"""
|
|
|
|
import os
|
|
import pickle # 用于数据序列化
|
|
import numpy as np
|
|
import tempfile
|
|
import trimesh
|
|
|
|
# 导入OpenCASCADE相关库
|
|
from OCC.Core.STEPControl import STEPControl_Reader # STEP文件读取器
|
|
from OCC.Core.TopExp import TopExp_Explorer # 拓扑结构遍历
|
|
from OCC.Core.TopAbs import TopAbs_FACE, TopAbs_EDGE, TopAbs_VERTEX # 拓扑类型定义
|
|
from OCC.Core.BRep import BRep_Tool # B-rep工具
|
|
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh # 网格剖分
|
|
from OCC.Core.TopLoc import TopLoc_Location # 位置变换
|
|
from OCC.Core.IFSelect import IFSelect_RetDone,IFSelect_RetError, IFSelect_RetFail, IFSelect_RetVoid # 操作状态码
|
|
from OCC.Core.TopoDS import topods # 拓扑数据结构
|
|
from OCC.Core.StlAPI import StlAPI_Writer
|
|
from OCC.Core.BRepAdaptor import BRepAdaptor_Surface
|
|
from OCC.Core.gp import gp_Pnt, gp_Vec
|
|
|
|
from brep2sdf.data.sampler import sample_sdf_points_and_normals, sample_face_points_brep, sample_edge_points_brep,sample_zero_surface_points_and_normals
|
|
from brep2sdf.data.data import check_data_format
|
|
from brep2sdf.data.utils import get_bbox, normalize, get_adjacency_info,batch_compute_normals
|
|
from brep2sdf.utils.logger import logger
|
|
# 导入配置
|
|
from brep2sdf.config.default_config import get_default_config
|
|
config = get_default_config()
|
|
|
|
# 设置最大面数阈值，用于加速处理
|
|
MAX_FACE = config.data.max_face
|
|
|
|
|
|
def parse_solid(step_path,sample_normal_vector=False,sample_sdf_points=False):
|
|
"""
|
|
解析STEP文件中的CAD模型数据
|
|
|
|
返回:
|
|
dict: 包含以下键值对的字典:
|
|
# 几何数据
|
|
'train_surf_ncs' np.ndarray(dtype=object) # 形状为(N,)的数组，每个元素是形状为(M, 3)的float32数组，表示归一化后的面点云
|
|
'surf_wcs': np.ndarray(dtype=object) # 形状为(N,)的数组，每个元素是形状为(M, 3)的float32数组，表示面的点云坐标，很多是边缘点，不合适训练
|
|
'edge_wcs': np.ndarray(dtype=object) # 形状为(N,)的数组，每个元素是形状为(num_edge_sample_points, 3)的float32数组，表示边的采样点坐标
|
|
'surf_ncs': np.ndarray(dtype=object) # 形状为(N,)的数组，每个元素是形状为(M, 3)的float32数组，表示归一化后的面点云
|
|
'edge_ncs': np.ndarray(dtype=object) # 形状为(N,)的数组，每个元素是形状为(num_edge_sample_points, 3)的float32数组，表示归一化后的边采样点
|
|
'corner_wcs': np.ndarray(dtype=float32) # 形状为(num_edges, 2, 3)的数组，表示每条边的两个端点坐标
|
|
'corner_unique': np.ndarray(dtype=float32) # 形状为(num_vertices, 3)的数组，表示所有顶点的唯一坐标,num_vertices <= num_edges * 2
|
|
|
|
# 拓扑关系
|
|
'edgeFace_adj': np.ndarray(dtype=int32) # 形状为(num_edges, num_faces)的数组，表示边-面邻接关系
|
|
'edgeCorner_adj': np.ndarray(dtype=int32) # 形状为(num_edges, 2)的数组，表示边-顶点邻接关系
|
|
'faceEdge_adj': np.ndarray(dtype=int32) # 形状为(num_faces, num_edges)的数组，表示面-边邻接关系
|
|
|
|
# 包围盒数据
|
|
'surf_bbox_wcs': np.ndarray(dtype=float32) # 形状为(num_faces, 6)的数组，表示每个面的包围盒[xmin,ymin,zmin,xmax,ymax,zmax]
|
|
'edge_bbox_wcs': np.ndarray(dtype=float32) # 形状为(num_edges, 6)的数组，表示每条边的包围盒[xmin,ymin,zmin,xmax,ymax,zmax]
|
|
"""
|
|
# Load STEP file
|
|
reader = STEPControl_Reader()
|
|
status = reader.ReadFile(step_path)
|
|
if status != IFSelect_RetDone:
|
|
if status == IFSelect_RetError:
|
|
print("Error: An error occurred while reading the file.")
|
|
elif status == IFSelect_RetFail:
|
|
print("Error: Failed to read the file.")
|
|
elif status == IFSelect_RetVoid:
|
|
print("Error: No data was read from the file.")
|
|
else:
|
|
print(f"Unexpected status code: {status}")
|
|
raise Exception(f"Failed to read STEP file. {status}")
|
|
|
|
reader.TransferRoots()
|
|
shape = reader.OneShape()
|
|
|
|
# Create mesh
|
|
mesh = BRepMesh_IncrementalMesh(shape, 0.01)
|
|
mesh.Perform()
|
|
|
|
# Initialize explorers
|
|
face_explorer = TopExp_Explorer(shape, TopAbs_FACE)
|
|
edge_explorer = TopExp_Explorer(shape, TopAbs_EDGE)
|
|
vertex_explorer = TopExp_Explorer(shape, TopAbs_VERTEX)
|
|
|
|
#tarin_surf_pnts = []
|
|
face_pnts = []
|
|
edge_pnts = []
|
|
corner_pnts = []
|
|
surf_bbox_wcs = []
|
|
edge_bbox_wcs = []
|
|
|
|
faces, edges, vertices = [], [], []
|
|
|
|
# Extract face points
|
|
logger.debug("Extract face points...")
|
|
total_sample_points = config.data.num_surf_points # 总采样点数
|
|
min_points_per_face = 50 # 每个面的最小采样点数
|
|
|
|
# 第一次遍历：收集所有面的原始点数和面积
|
|
face_areas = []
|
|
original_points = []
|
|
while face_explorer.More():
|
|
face = topods.Face(face_explorer.Current())
|
|
faces.append(face)
|
|
loc = TopLoc_Location()
|
|
triangulation = BRep_Tool.Triangulation(face, loc)
|
|
|
|
points = []
|
|
if triangulation is not None:
|
|
for i in range(1, triangulation.NbNodes() + 1):
|
|
node = triangulation.Node(i)
|
|
pnt = node.Transformed(loc.Transformation())
|
|
points.append([pnt.X(), pnt.Y(), pnt.Z()])
|
|
|
|
if points:
|
|
points = np.array(points, dtype=np.float32)
|
|
if len(points.shape) == 2 and points.shape[1] == 3:
|
|
original_points.append(points)
|
|
# 计算面积
|
|
surface = BRepAdaptor_Surface(face)
|
|
u_min, u_max = surface.FirstUParameter(), surface.LastUParameter()
|
|
v_min, v_max = surface.FirstVParameter(), surface.LastVParameter()
|
|
face_areas.append((u_max - u_min) * (v_max - v_min)) # 简化的面积计算
|
|
face_explorer.Next()
|
|
|
|
face_count = len(faces)
|
|
if face_count > MAX_FACE:
|
|
logger.error(f"step has {face_count} faces, which exceeds MAX_FACE {MAX_FACE}")
|
|
return None
|
|
|
|
# 计算每个面应该分配的采样点数
|
|
face_areas = np.array(face_areas)
|
|
total_area = np.sum(face_areas)
|
|
# 根据面积分配点数，但确保每个面至少有min_points_per_face个点
|
|
points_per_face = np.maximum(
|
|
min_points_per_face,
|
|
(face_areas / total_area * (total_sample_points - min_points_per_face * face_count)).astype(int) + min_points_per_face
|
|
)
|
|
|
|
# 第二次遍历：对每个面进行采样
|
|
face_explorer = TopExp_Explorer(shape, TopAbs_FACE) # 重置explorer
|
|
face_idx = 0
|
|
|
|
while face_explorer.More():
|
|
face = topods.Face(face_explorer.Current())
|
|
if face_idx < len(original_points):
|
|
points = original_points[face_idx]
|
|
target_points = points_per_face[face_idx]
|
|
#tarin_surf_pnts.append(sample_face_points_brep(face, min_points=target_points))
|
|
# 如果需要补充采样
|
|
if len(points) < target_points:
|
|
try:
|
|
# 使用sample_face_points_brep补充采样
|
|
extra_points_needed = target_points - len(points)
|
|
extra_points = sample_face_points_brep(face, min_points=extra_points_needed)
|
|
|
|
if extra_points is not None:
|
|
points = np.vstack([points, extra_points])
|
|
|
|
|
|
except Exception as e:
|
|
logger.warning(f"补充采样点失败: {str(e)}")
|
|
|
|
# 如果点数超过目标，随机采样到目标点数
|
|
if len(points) > target_points:
|
|
indices = np.random.choice(len(points), target_points, replace=False)
|
|
points = points[indices]
|
|
|
|
face_pnts.append(points)
|
|
surf_bbox_wcs.append(get_bbox(shape, face))
|
|
face_idx += 1
|
|
|
|
face_explorer.Next()
|
|
|
|
# Extract edge points
|
|
logger.debug("Extract edge points...")
|
|
num_samples = config.model.num_edge_points # 使用配置中的边采样点数
|
|
while edge_explorer.More():
|
|
edge = topods.Edge(edge_explorer.Current())
|
|
edges.append(edge)
|
|
#logger.debug(len(edges))
|
|
curve_info = BRep_Tool.Curve(edge)
|
|
if curve_info is None:
|
|
continue # 跳过无效边
|
|
|
|
try:
|
|
if len(curve_info) == 3:
|
|
curve, first, last = curve_info
|
|
elif len(curve_info) == 2:
|
|
curve = None # 跳过判断
|
|
else:
|
|
raise ValueError(f"Unexpected curve info: {curve_info}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to process edge {edge}: {str(e)}")
|
|
curve = None
|
|
|
|
if curve is not None:
|
|
points = []
|
|
for i in range(num_samples):
|
|
param = first + (last - first) * float(i) / (num_samples - 1)
|
|
pnt = curve.Value(param)
|
|
points.append([pnt.X(), pnt.Y(), pnt.Z()])
|
|
|
|
if points:
|
|
points = np.array(points, dtype=np.float32)
|
|
if len(points.shape) == 2 and points.shape[1] == 3:
|
|
edge_pnts.append(points) # 现在points是(num_edge_points, 3)形状
|
|
edge_bbox_wcs.append(get_bbox(shape, edge))
|
|
|
|
edge_explorer.Next()
|
|
|
|
# Extract vertex points
|
|
logger.debug("Extract vertex points...")
|
|
while vertex_explorer.More():
|
|
vertex = topods.Vertex(vertex_explorer.Current())
|
|
vertices.append(vertex)
|
|
pnt = BRep_Tool.Pnt(vertex)
|
|
corner_pnts.append([pnt.X(), pnt.Y(), pnt.Z()])
|
|
vertex_explorer.Next()
|
|
|
|
# 获取邻接信息
|
|
edgeFace_adj, faceEdge_adj, edgeCorner_adj = get_adjacency_info(
|
|
shape,
|
|
faces=faces, # 传入已收集的面列表
|
|
edges=edges, # 传入已收集的边列表
|
|
vertices=vertices # 传入已收集的顶点列表
|
|
)
|
|
logger.debug("complete.")
|
|
|
|
# 转换为numpy数组时确保类型正确
|
|
face_pnts = [np.array(points, dtype=np.float32) for points in face_pnts]
|
|
edge_pnts = [np.array(points, dtype=np.float32) for points in edge_pnts]
|
|
|
|
# 转换为对象数组
|
|
face_pnts = np.array(face_pnts, dtype=object)
|
|
edge_pnts = np.array(edge_pnts, dtype=object)
|
|
corner_pnts = np.array(corner_pnts, dtype=np.float32)
|
|
|
|
# 重组顶点数据为每条边两个端点的形式
|
|
corner_pairs = []
|
|
for edge_idx in range(len(edge_pnts)):
|
|
v1_idx, v2_idx = edgeCorner_adj[edge_idx]
|
|
v1_pos = corner_pnts[v1_idx]
|
|
v2_pos = corner_pnts[v2_idx]
|
|
# 按坐标排序确保顺序一致
|
|
if (v1_pos > v2_pos).any():
|
|
v1_pos, v2_pos = v2_pos, v1_pos
|
|
corner_pairs.append(np.stack([v1_pos, v2_pos]))
|
|
|
|
corner_pairs = np.stack(corner_pairs).astype(np.float32) # [num_edges, 2, 3]
|
|
|
|
# 确保所有数组都有正确的类型
|
|
surf_bbox_wcs = np.array(surf_bbox_wcs, dtype=np.float32)
|
|
edge_bbox_wcs = np.array(edge_bbox_wcs, dtype=np.float32)
|
|
|
|
# Normalize the CAD model
|
|
surfs_wcs, edges_wcs, surfs_ncs, edges_ncs, corner_wcs,center, scale = normalize(
|
|
face_pnts, edge_pnts, corner_pairs)
|
|
|
|
# 计算归一化后的包围盒
|
|
surf_bbox_ncs = np.empty_like(surf_bbox_wcs)
|
|
edge_bbox_ncs = np.empty_like(edge_bbox_wcs)
|
|
|
|
# 转换曲面包围盒到归一化坐标系
|
|
surf_bbox_ncs[:, :3] = (surf_bbox_wcs[:, :3] - center) * scale # 最小点
|
|
surf_bbox_ncs[:, 3:] = (surf_bbox_wcs[:, 3:] - center) * scale # 最大点
|
|
|
|
# 转换边包围盒到归一化坐标系
|
|
edge_bbox_ncs[:, :3] = (edge_bbox_wcs[:, :3] - center) * scale # 最小点
|
|
edge_bbox_ncs[:, 3:] = (edge_bbox_wcs[:, 3:] - center) * scale # 最大点
|
|
|
|
|
|
# 验证归一化后的数据
|
|
if any(x is None for x in [surfs_wcs, edges_wcs, surfs_ncs, edges_ncs, corner_wcs]):
|
|
logger.error(f"Normalization failed for {step_path}")
|
|
return None
|
|
# --- 计算边的类型 ---
|
|
logger.debug("计算边的类型...")
|
|
edge_types = [] # 0:凹边 1:凸边
|
|
|
|
# 对每条边进行处理
|
|
for edge_idx in range(len(edges)):
|
|
# 获取与当前边相邻的面
|
|
adjacent_faces = np.where(edgeFace_adj[edge_idx] == 1)[0]
|
|
|
|
# 如果边只有一个相邻面或没有相邻面，默认为凹边
|
|
if len(adjacent_faces) < 2:
|
|
edge_types.append(0)
|
|
continue
|
|
|
|
# 获取两个相邻面
|
|
face1, face2 = faces[adjacent_faces[0]], faces[adjacent_faces[1]]
|
|
|
|
# 获取第一个面的法向量
|
|
surf1 = BRepAdaptor_Surface(face1)
|
|
u1 = (surf1.FirstUParameter() + surf1.LastUParameter()) / 2
|
|
v1 = (surf1.FirstVParameter() + surf1.LastVParameter()) / 2
|
|
pnt1 = gp_Pnt()
|
|
du1 = gp_Vec()
|
|
dv1 = gp_Vec()
|
|
surf1.D1(u1, v1, pnt1, du1, dv1)
|
|
normal1 = du1.Crossed(dv1)
|
|
normal1.Normalize()
|
|
normal1_np = np.array([normal1.X(), normal1.Y(), normal1.Z()])
|
|
|
|
# 获取第二个面的法向量
|
|
surf2 = BRepAdaptor_Surface(face2)
|
|
u2 = (surf2.FirstUParameter() + surf2.LastUParameter()) / 2
|
|
v2 = (surf2.FirstVParameter() + surf2.LastVParameter()) / 2
|
|
pnt2 = gp_Pnt()
|
|
du2 = gp_Vec()
|
|
dv2 = gp_Vec()
|
|
surf2.D1(u2, v2, pnt2, du2, dv2)
|
|
normal2 = du2.Crossed(dv2)
|
|
normal2.Normalize()
|
|
normal2_np = np.array([normal2.X(), normal2.Y(), normal2.Z()])
|
|
|
|
# 获取边的方向向量
|
|
edge = edges[edge_idx]
|
|
curve_info = BRep_Tool.Curve(edge)
|
|
if curve_info is None or len(curve_info) < 3:
|
|
edge_types.append(0)
|
|
continue
|
|
|
|
curve, first, last = curve_info
|
|
# 计算边的方向向量
|
|
start_point = np.array([curve.Value(first).X(), curve.Value(first).Y(), curve.Value(first).Z()])
|
|
end_point = np.array([curve.Value(last).X(), curve.Value(last).Y(), curve.Value(last).Z()])
|
|
edge_vector = end_point - start_point
|
|
edge_vector = edge_vector / np.linalg.norm(edge_vector)
|
|
|
|
# 使用混合积判断凹凸性
|
|
# 如果混合积为正，说明是凸边；为负，说明是凹边
|
|
mixed_product = np.dot(np.cross(normal1_np, normal2_np), edge_vector)
|
|
|
|
# 根据混合积的符号确定边的类型
|
|
edge_types.append(1 if mixed_product > 0 else 0)
|
|
|
|
edge_types = np.array(edge_types, dtype=np.int32)
|
|
|
|
# 创建结果字典并确保所有数组都有正确的类型
|
|
data = {
|
|
#'train_surf_ncs': np.array(train_surf_ncs, dtype=object), # 保持对象数组
|
|
'surf_wcs': np.array(surfs_wcs, dtype=object), # 保持对象数组
|
|
'edge_wcs': np.array(edges_wcs, dtype=object), # 保持对象数组
|
|
'surf_ncs': np.array(surfs_ncs, dtype=object), # 保持对象数组
|
|
'edge_ncs': np.array(edges_ncs, dtype=object), # 保持对象数组
|
|
'corner_wcs': corner_wcs.astype(np.float32), # [num_edges, 2, 3]
|
|
'edgeFace_adj': edgeFace_adj.astype(np.int32), # [num_edges, num_faces]， 1 表示边与面相邻
|
|
'edgeCorner_adj': edgeCorner_adj.astype(np.int32),
|
|
'faceEdge_adj': faceEdge_adj.astype(np.int32),
|
|
'edge_types': np.array(edge_types, dtype=np.int32), # [num_edges]
|
|
'surf_bbox_wcs': surf_bbox_wcs.astype(np.float32),
|
|
'edge_bbox_wcs': edge_bbox_wcs.astype(np.float32),
|
|
'surf_bbox_ncs': surf_bbox_ncs.astype(np.float32), # 归一化坐标系 [num_faces, 6]
|
|
'edge_bbox_ncs': edge_bbox_ncs.astype(np.float32), # 归一化坐标系 [num_edges, 6]
|
|
'corner_unique': np.unique(corner_wcs.reshape(-1, 3), axis=0).astype(np.float32), # 先展平再去重
|
|
'normalization_params': {
|
|
'center': center.astype(np.float32), # 归一化中心点 [3,]
|
|
'scale': float(scale), # 归一化缩放系数
|
|
}
|
|
}
|
|
|
|
trimesh_mesh = None
|
|
trimesh_mesh_ncs = None
|
|
|
|
# --- Trimesh 加载和处理 (如果需要) ---
|
|
if sample_normal_vector or sample_sdf_points:
|
|
logger.debug("加载 Trimesh 用于法线/SDF 采样...")
|
|
# 注意：这里的 mesh (BRepMesh_IncrementalMesh) 与 trimesh 不同
|
|
# 需要从原始 shape 导出 STL
|
|
stl_writer = StlAPI_Writer()
|
|
stl_writer.SetASCIIMode(False)
|
|
tmp_stl_path = ""
|
|
try:
|
|
with tempfile.NamedTemporaryFile(suffix='.stl', delete=True) as tmp:
|
|
tmp_stl_path = tmp.name
|
|
# 检查 shape 是否有效
|
|
if shape.IsNull():
|
|
raise ValueError("OCC Shape is Null, cannot write STL.")
|
|
success = stl_writer.Write(shape, tmp_stl_path)
|
|
if not success:
|
|
raise RuntimeError(f"StlAPI_Writer failed to write {tmp_stl_path}")
|
|
|
|
trimesh_mesh = trimesh.load(tmp_stl_path)
|
|
|
|
# 创建归一化 Trimesh
|
|
vertices_wcs = trimesh_mesh.vertices.astype(np.float32)
|
|
logger.debug(f"vertices_wcs:{vertices_wcs}")
|
|
logger.debug(f"center:{data['normalization_params']['center']},scale:{data['normalization_params']['scale']}")
|
|
vertices_ncs = (vertices_wcs - data['normalization_params']['center']) * data['normalization_params']['scale']
|
|
logger.debug(f"vertices_ncs:{vertices_ncs}")
|
|
trimesh_mesh_ncs = trimesh.Trimesh(vertices=vertices_ncs, faces=trimesh_mesh.faces, process=False)
|
|
|
|
if not trimesh_mesh_ncs.is_watertight:
|
|
logger.debug(f"{step_path} 的归一化网格不是 watertight，尝试修复。")
|
|
trimesh.repair.fill_holes(trimesh_mesh_ncs)
|
|
if not trimesh_mesh_ncs.is_watertight:
|
|
logger.warning(f"{step_path} 的归一化网格修复后仍不是 watertight。")
|
|
|
|
data["train_surf_ncs"] = sample_zero_surface_points_and_normals(trimesh_mesh_ncs, config.data.num_surf_points) # 归一化网格的顶点
|
|
|
|
except Exception as e:
|
|
logger.error(f"为 {step_path} 加载/处理 Trimesh 失败: {e}")
|
|
trimesh_mesh = None
|
|
trimesh_mesh_ncs = None
|
|
# 如果你需要归一化后的表面点
|
|
# --- 计算表面点法线 ---
|
|
if sample_normal_vector and trimesh_mesh_ncs is not None:
|
|
logger.debug("计算表面点法线...")
|
|
# 使用 data['surf_ncs'] 因为它们已经是归一化后的点云
|
|
if data['surf_ncs'].shape[0] > 0:
|
|
# 确保 batch_compute_normals 使用归一化的 mesh
|
|
data['surf_pnt_normals'] = batch_compute_normals(trimesh_mesh_ncs, data['surf_ncs'])
|
|
else:
|
|
logger.warning("没有有效的归一化表面点云用于法线计算。")
|
|
data['surf_pnt_normals'] = np.array([], dtype=object)
|
|
elif sample_normal_vector:
|
|
logger.warning("请求了表面法线计算，但 Trimesh 加载失败。")
|
|
data['surf_pnt_normals'] = np.array([], dtype=object) # 添加空键
|
|
|
|
# --- SDF 点采样 ---
|
|
data['sampled_points_normals_sdf'] = None # 初始化键
|
|
if sample_sdf_points:
|
|
if trimesh_mesh_ncs is not None:
|
|
# 调用封装的函数，传递固定数量参数
|
|
logger.debug("采样 SDF 点和法线...")
|
|
data['sampled_points_normals_sdf'] = sample_sdf_points_and_normals(
|
|
trimesh_mesh_ncs=trimesh_mesh_ncs,
|
|
surf_bbox_ncs=data['surf_bbox_ncs'],
|
|
num_sdf_samples=50000, # <-- 传递固定数量
|
|
sdf_sampling_std_dev=0.0001
|
|
)
|
|
else:
|
|
logger.warning("请求了 SDF 点采样，但 Trimesh 加载失败。")
|
|
return data
|
|
|
|
def process_single_step(step_path:str, output_path:str=None, sample_normal_vector=False, sample_sdf_points=False, timeout:int=300) -> dict:
|
|
"""处理单个STEP文件, 从 brep 2 pkl
|
|
return data = {
|
|
'train_surf_ncs' np.ndarray(dtype=object) # 形状为(N,)的数组，每个元素是形状为(M, 3)的float32数组，表示归一化后的面点云
|
|
'surf_wcs': np.array(surfs_wcs, dtype=object), # 世界坐标系下的曲面几何数据(对象数组)
|
|
'edge_wcs': np.array(edges_wcs, dtype=object), # 世界坐标系下的边几何数据(对象数组)
|
|
'surf_ncs': np.array(surfs_ncs, dtype=object), # 归一化坐标系下的曲面几何数据(对象数组) 面归一化点云 [num_faces, num_surf_sample_points, 3]
|
|
'edge_ncs': np.array(edges_ncs, dtype=object), # 归一化坐标系下的边几何数据(对象数组) 边归一化点云 [num_edges, num_edge_sample_points, 3]
|
|
'corner_wcs': corner_wcs.astype(np.float32), # 世界坐标系下的角点数据 [num_edges, 2, 3]
|
|
'edgeFace_adj': edgeFace_adj.astype(np.int32), # 边-面的邻接关系矩阵
|
|
'edgeCorner_adj': edgeCorner_adj.astype(np.int32),# 边-角点的邻接关系矩阵
|
|
'faceEdge_adj': faceEdge_adj.astype(np.int32), # 面-边的邻接关系矩阵
|
|
'edge_types': np.array(edge_types, dtype=np.int32)# [num_edges]
|
|
'surf_bbox_wcs': surf_bbox_wcs.astype(np.float32),# 曲面在世界坐标系下的包围盒
|
|
'edge_bbox_wcs': edge_bbox_wcs.astype(np.float32),# 边在世界坐标系下的包围盒
|
|
'corner_unique': np.unique(corner_wcs.reshape(-1, 3), axis=0).astype(np.float32) # 去重后的唯一角点坐标
|
|
'normalization_params': { # 归一化参数
|
|
'center': center.astype(np.float32), # 归一化中心点 [3,]
|
|
'scale': float(scale), # 归一化缩放系数
|
|
},
|
|
'surf_pnt_normals': np.array(dtype=object), # 表面点的法线数据 [num_faces, num_surf_sample_points, 3]，仅当 sample_normal_vector=True
|
|
'sampled_points_normals_sdf': np.array(dtype=float32), # 采样点的位置、法线和SDF值 [num_samples, 7]，仅当 sample_sdf_points=True
|
|
}"""
|
|
try:
|
|
logger.info("数据预处理……")
|
|
if not os.path.exists(step_path):
|
|
logger.error(f"STEP文件不存在: {step_path}")
|
|
return None
|
|
if not step_path.lower().endswith('.step') and not step_path.lower().endswith('.stp'):
|
|
logger.error(f"文件格式不支持，必须是.step或.stp文件: {step_path}")
|
|
return None
|
|
# 解析STEP文件
|
|
data = parse_solid(step_path, sample_normal_vector,sample_sdf_points)
|
|
if data is None:
|
|
logger.error(f"Failed to parse STEP file: {step_path}")
|
|
return None
|
|
|
|
# 检查数据格式
|
|
is_valid, msg = check_data_format(data, step_path)
|
|
if not is_valid:
|
|
logger.error(f"Data format check failed for {step_path}: {msg}")
|
|
return None
|
|
|
|
# 保存结果
|
|
if output_path:
|
|
try:
|
|
logger.debug(f"Saving results to: {output_path}")
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
with open(output_path, 'wb') as f:
|
|
pickle.dump(data, f)
|
|
logger.debug(f"Results saved successfully: {output_path}")
|
|
return data
|
|
except Exception as e:
|
|
logger.error(f'Failed to save {output_path}: {str(e)}')
|
|
return None
|
|
logger.info("数据预处理完成")
|
|
return data
|
|
|
|
except Exception as e:
|
|
logger.error(f'Error processing {step_path}: {str(e)}')
|
|
return None
|
|
|
|
def test(step_file_path, output_path=None):
|
|
"""
|
|
测试函数：转换单个STEP文件并保存结果
|
|
"""
|
|
try:
|
|
logger.info(f"Processing STEP file: {step_file_path}")
|
|
|
|
# 解析STEP文件
|
|
data = parse_solid(step_file_path)
|
|
if data is None:
|
|
logger.error(f"Failed to parse STEP file: {step_file_path}")
|
|
return None
|
|
|
|
# 检查数据格式
|
|
is_valid, msg = check_data_format(data, step_file_path)
|
|
if not is_valid:
|
|
logger.error(f"Data format check failed for {step_file_path}: {msg}")
|
|
return None
|
|
|
|
# 打印统计信息
|
|
logger.debug("\nStatistics:")
|
|
logger.debug(f"Number of surfaces: {len(data['surf_wcs'])}")
|
|
logger.debug(f"Number of edges: {len(data['edge_wcs'])}")
|
|
logger.debug(f"Number of corners: {len(data['corner_wcs'])}") # 修正为corner_wcs
|
|
|
|
# 保存结果
|
|
if output_path:
|
|
try:
|
|
logger.debug(f"Saving results to: {output_path}")
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
with open(output_path, 'wb') as f:
|
|
pickle.dump(data, f)
|
|
logger.debug(f"Results saved successfully: {output_path}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to save {output_path}: {str(e)}")
|
|
return None
|
|
|
|
return data
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing {step_file_path}: {str(e)}")
|
|
return None
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# main()
|
|
test("/home/wch/brep2sdf/data/step/00000000/00000000_290a9120f9f249a7a05cfe9c_step_000.step","/home/wch/brep2sdf/test_data/pkl/train/00000031xx.pkl")
|
|
#test("/home/wch/brep2sdf/00000031_ad34a3f60c4a4caa99646600_step_011.step", "/home/wch/brep2sdf/test_data/pkl/train/00000031.pkl")
|
|
#test("/mnt/mynewdisk/dataset/furniture/step/furniture_dataset_step/train/bathtub_0004.step", "/home/wch/brep2sdf/test_data/pkl/train/0004.pkl")
|
|
#reader = STEPControl_Reader()
|
|
|