""" CAD模型处理脚本 功能：将STEP格式的CAD模型转换为结构化数据，包括： - 几何信息：面、边、顶点的坐标数据 - 拓扑信息：面-边-顶点的邻接关系 - 空间信息：包围盒数据 """ import os import pickle # 用于数据序列化 import argparse # 命令行参数解析 import numpy as np from tqdm import tqdm # 进度条显示 from concurrent.futures import ProcessPoolExecutor, as_completed, TimeoutError # 并行处理 import logging from datetime import datetime # 创建logs目录 os.makedirs('logs', exist_ok=True) # 设置日志记录器 logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # 创建格式化器 formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') # 创建文件处理器 current_time = datetime.now().strftime('%Y%m%d_%H%M%S') log_file = f'logs/process_brep_{current_time}.log' file_handler = logging.FileHandler(log_file, encoding='utf-8') file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) # 添加文件处理器到日志记录器 logger.addHandler(file_handler) # 记录脚本开始执行 logger.info("="*50) logger.info("Script started") logger.info(f"Log file: {log_file}") logger.info("="*50) # 导入OpenCASCADE相关库 from OCC.Core.STEPControl import STEPControl_Reader # STEP文件读取器 from OCC.Core.TopExp import TopExp_Explorer, topexp # 拓扑结构遍历 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 # 操作状态码 from OCC.Core.TopTools import TopTools_IndexedDataMapOfShapeListOfShape # 形状映射 from OCC.Core.BRepBndLib import brepbndlib # 包围盒计算 from OCC.Core.Bnd import Bnd_Box # 包围盒 from OCC.Core.TopoDS import TopoDS_Shape, topods, TopoDS_Vertex # 拓扑数据结构 # 设置最大面数阈值，用于加速处理 MAX_FACE = 70 def normalize(surfs, edges, corners): """ 将CAD模型归一化到单位立方体空间 参数: surfs: 面的点集列表 edges: 边的点集列表 corners: 顶点坐标列表 返回: surfs_wcs: 原始坐标系下的面点集 edges_wcs: 原始坐标系下的边点集 surfs_ncs: 归一化坐标系下的面点集 edges_ncs: 归一化坐标系下的边点集 corner_wcs: 归一化后的顶点坐标 """ if len(corners) == 0: return None, None, None, None, None # 计算包围盒和缩放因子 corners_array = np.array(corners) center = (corners_array.max(0) + corners_array.min(0)) / 2 # 计算中心点 scale = 1.0 / (corners_array.max(0) - corners_array.min(0)).max() # 计算缩放系数 # 归一化面的点集 surfs_wcs = [] # 原始世界坐标系下的面点集 surfs_ncs = [] # 归一化坐标系下的面点集 for surf in surfs: surf_wcs = np.array(surf) surf_ncs = (surf_wcs - center) * scale # 归一化变换 surfs_wcs.append(surf_wcs) surfs_ncs.append(surf_ncs) # 归一化边的点集 edges_wcs = [] # 原始世界坐标系下的边点集 edges_ncs = [] # 归一化坐标系下的边点集 for edge in edges: edge_wcs = np.array(edge) edge_ncs = (edge_wcs - center) * scale # 归一化变换 edges_wcs.append(edge_wcs) edges_ncs.append(edge_ncs) # 归一化顶点坐标 corner_wcs = (corners_array - center) * scale return (np.array(surfs_wcs, dtype=object), np.array(edges_wcs, dtype=object), np.array(surfs_ncs, dtype=object), np.array(edges_ncs, dtype=object), corner_wcs) def get_adjacency_info(shape): """ 获取CAD模型中面、边、顶点之间的邻接关系 参数: shape: CAD模型的形状对象 返回: edgeFace_adj: 边-面邻接矩阵 (num_edges × num_faces) faceEdge_adj: 面-边邻接矩阵 (num_faces × num_edges) edgeCorner_adj: 边-顶点邻接矩阵 (num_edges × 2) """ # 创建边-面映射关系 edge_face_map = TopTools_IndexedDataMapOfShapeListOfShape() topexp.MapShapesAndAncestors(shape, TopAbs_EDGE, TopAbs_FACE, edge_face_map) # 获取所有几何元素 faces = [] # 存储所有面 edges = [] # 存储所有边 vertices = [] # 存储所有顶点 # 创建拓扑结构探索器 face_explorer = TopExp_Explorer(shape, TopAbs_FACE) edge_explorer = TopExp_Explorer(shape, TopAbs_EDGE) vertex_explorer = TopExp_Explorer(shape, TopAbs_VERTEX) # 收集所有几何元素 while face_explorer.More(): faces.append(topods.Face(face_explorer.Current())) face_explorer.Next() while edge_explorer.More(): edges.append(topods.Edge(edge_explorer.Current())) edge_explorer.Next() while vertex_explorer.More(): vertices.append(topods.Vertex(vertex_explorer.Current())) vertex_explorer.Next() # 创建邻接矩阵 num_faces = len(faces) num_edges = len(edges) num_vertices = len(vertices) edgeFace_adj = np.zeros((num_edges, num_faces), dtype=np.int32) faceEdge_adj = np.zeros((num_faces, num_edges), dtype=np.int32) edgeCorner_adj = np.zeros((num_edges, 2), dtype=np.int32) # 填充边-面邻接矩阵 for i, edge in enumerate(edges): # 检查每个面是否与当前边相连 for j, face in enumerate(faces): edge_explorer = TopExp_Explorer(face, TopAbs_EDGE) while edge_explorer.More(): if edge.IsSame(edge_explorer.Current()): edgeFace_adj[i, j] = 1 faceEdge_adj[j, i] = 1 break edge_explorer.Next() # 获取边的两个端点 v1 = TopoDS_Vertex() v2 = TopoDS_Vertex() topexp.Vertices(edge, v1, v2) # 记录边的端点索引 if not v1.IsNull() and not v2.IsNull(): v1_vertex = topods.Vertex(v1) v2_vertex = topods.Vertex(v2) for k, vertex in enumerate(vertices): if v1_vertex.IsSame(vertex): edgeCorner_adj[i, 0] = k if v2_vertex.IsSame(vertex): edgeCorner_adj[i, 1] = k return edgeFace_adj, faceEdge_adj, edgeCorner_adj def get_bbox(shape, subshape): """ 计算形状的包围盒 参数: shape: 完整的CAD模型形状 subshape: 需要计算包围盒的子形状（面或边） 返回: 包围盒的六个参数 [xmin, ymin, zmin, xmax, ymax, zmax] """ bbox = Bnd_Box() brepbndlib.Add(subshape, bbox) xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get() return np.array([xmin, ymin, zmin, xmax, ymax, zmax]) def parse_solid(step_path): """Parse the surface, curve, face, edge, vertex in a CAD solid using OCC.""" # Load STEP file reader = STEPControl_Reader() status = reader.ReadFile(step_path) if status != IFSelect_RetDone: raise Exception("Failed to read STEP file") 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) face_pnts = [] edge_pnts = [] corner_pnts = [] surf_bbox_wcs = [] edge_bbox_wcs = [] # Extract face points while face_explorer.More(): face = topods.Face(face_explorer.Current()) loc = TopLoc_Location() triangulation = BRep_Tool.Triangulation(face, loc) if triangulation is not None: points = [] 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: face_pnts.append(points) surf_bbox_wcs.append(get_bbox(shape, face)) face_explorer.Next() # Extract edge points num_samples = 100 while edge_explorer.More(): edge = topods.Edge(edge_explorer.Current()) curve, first, last = BRep_Tool.Curve(edge) 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) edge_bbox_wcs.append(get_bbox(shape, edge)) edge_explorer.Next() # Extract vertex points while vertex_explorer.More(): vertex = topods.Vertex(vertex_explorer.Current()) 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) # 转换为numpy数组，但保持列表形式 face_pnts = list(face_pnts) # 确保是列表 edge_pnts = list(edge_pnts) # 确保是列表 corner_pnts = np.array(corner_pnts, dtype=np.float32) 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 = normalize(face_pnts, edge_pnts, corner_pnts) # Create result dictionary data = { 'surf_wcs': surfs_wcs, 'edge_wcs': edges_wcs, 'surf_ncs': surfs_ncs, 'edge_ncs': edges_ncs, 'corner_wcs': corner_wcs.astype(np.float32), 'edgeFace_adj': edgeFace_adj, 'edgeCorner_adj': edgeCorner_adj, 'faceEdge_adj': faceEdge_adj, 'surf_bbox_wcs': surf_bbox_wcs, 'edge_bbox_wcs': edge_bbox_wcs, 'corner_unique': np.unique(corner_wcs, axis=0).astype(np.float32) } return data def load_step(step_path): """Load STEP file and return solids""" reader = STEPControl_Reader() reader.ReadFile(step_path) reader.TransferRoots() return [reader.OneShape()] def process_single_step( step_path:str, output_path:str=None, timeout:int=300 ) -> dict: """Process single STEP file""" try: # 解析STEP文件 data = parse_solid(step_path) if data is None: logger.error("Failed to parse STEP file") return None # 保存结果 if output_path: try: logger.info(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.info("Results saved successfully") except Exception as e: logger.error(f'Not saving due to error: {str(e)}') return data except Exception as e: logger.error(f'Not saving due to error: {str(e)}') return 0 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("Failed to parse STEP file") return None # 打印统计信息 logger.info("\nStatistics:") logger.info(f"Number of surfaces: {len(data['surf_wcs'])}") logger.info(f"Number of edges: {len(data['edge_wcs'])}") logger.info(f"Number of corners: {len(data['corner_unique'])}") # 保存结果 if output_path: logger.info(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.info("Results saved successfully") return data except Exception as e: logger.error(f"Error processing STEP file: {str(e)}") return None def process_furniture_step(data_path): """ 处理家具数据集的STEP文件 参数: data_path: 数据集路径 返回: 包含训练、验证和测试集的STEP文件路径字典 { 'train': [step_file_path1, step_file_path2, ...], 'val': [step_file_path1, step_file_path2, ...], 'test': [step_file_path1, step_file_path2, ...] } """ step_dirs = {} for split in ['train', 'val', 'test']: tmp_step_dirs = [] split_path = os.path.join(data_path, split) if os.path.exists(split_path): for f in os.listdir(split_path): if f.endswith('.step'): tmp_step_dirs.append(f) step_dirs[split] = tmp_step_dirs return step_dirs def main(): """ 主函数：处理多个STEP文件 """ # 定义路径常量 INPUT = '/mnt/disk2/dataset/furniture/step/furniture_dataset_step/' OUTPUT = 'test_data/pkl/' RESULT = 'test_data/result/' # 用于存储成功/失败文件记录 # 确保输出目录存在 os.makedirs(OUTPUT, exist_ok=True) os.makedirs(RESULT, exist_ok=True) # 获取所有STEP文件 step_dirs_dict = process_furniture_step(INPUT) total_processed = 0 total_success = 0 # 按数据集分割处理文件 for split in ['train', 'val', 'test']: current_step_dirs = step_dirs_dict[split] if not current_step_dirs: logger.warning(f"No files found in {split} split") continue # 确保分割目录存在 split_output_dir = os.path.join(OUTPUT, split) split_result_dir = os.path.join(RESULT, split) os.makedirs(split_output_dir, exist_ok=True) os.makedirs(split_result_dir, exist_ok=True) success_files = [] # 存储成功处理的文件名 failed_files = [] # 存储失败的文件名及原因 # 并行处理文件 with ProcessPoolExecutor(max_workers=os.cpu_count() // 2) as executor: futures = {} for step_file in current_step_dirs: input_path = os.path.join(INPUT, split, step_file) output_path = os.path.join(split_output_dir, step_file.replace('.step', '.pkl')) future = executor.submit(process_single_step, input_path, output_path, timeout=300) futures[future] = step_file # 处理结果 for future in tqdm(as_completed(futures), total=len(current_step_dirs), desc=f"Processing {split} set"): try: status = future.result(timeout=300) if status is not None: success_files.append(futures[future]) total_success += 1 except TimeoutError: logger.error(f"Timeout occurred while processing {futures[future]}") failed_files.append((futures[future], "Timeout")) except Exception as e: logger.error(f"Error processing {futures[future]}: {str(e)}") failed_files.append((futures[future], str(e))) finally: total_processed += 1 # 保存成功文件列表 success_file_path = os.path.join(split_result_dir, 'success.txt') with open(success_file_path, 'w', encoding='utf-8') as f: f.write('\n'.join(success_files)) logger.info(f"Saved {len(success_files)} successful files to {success_file_path}") # 保存失败文件列表（包含错误信息） failed_file_path = os.path.join(split_result_dir, 'failed.txt') with open(failed_file_path, 'w', encoding='utf-8') as f: for file, error in failed_files: f.write(f"{file}: {error}\n") logger.info(f"Saved {len(failed_files)} failed files to {failed_file_path}") # 打印最终统计信息 if total_processed > 0: success_rate = (total_success / total_processed) * 100 logger.info(f"Processing completed:") logger.info(f"Total files processed: {total_processed}") logger.info(f"Successfully processed: {total_success}") logger.info(f"Success rate: {success_rate:.2f}%") else: logger.warning("No files were processed") if __name__ == '__main__': main()