diff --git a/brep2sdf/data/utils.py b/brep2sdf/data/utils.py index f847303..cfc31bd 100644 --- a/brep2sdf/data/utils.py +++ b/brep2sdf/data/utils.py @@ -1091,10 +1091,25 @@ def process_brep_data( random_indices = np.random.permutation(surf_pos.shape[0]) surf_pos = surf_pos[random_indices] # [num_faces, 6] edge_pos = edge_pos[random_indices] # [num_faces, max_edge, 6] - # 修改这里：surf_ncs是对象数组，需要特殊处理 - surf_ncs_new = np.array([surf_ncs[i] for i in random_indices]) # 重新排列对象数组 - surf_ncs = surf_ncs_new - #surf_ncs = surf_ncs[random_indices] # [num_faces, 100, 3] + + # 处理surf_ncs (对象数组) + surf_ncs_list = [] + for idx in random_indices: + points = surf_ncs[idx] # 获取当前面的点云 + # 确保点云数据形状正确 (N, 3) -> (100, 3) + if len(points) > 100: + # 如果点数超过100，随机采样 + indices = np.random.choice(len(points), 100, replace=False) + points = points[indices] + elif len(points) < 100: + # 如果点数少于100，重复采样 + indices = np.random.choice(len(points), 100-len(points)) + points = np.concatenate([points, points[indices]], axis=0) + surf_ncs_list.append(points) + + # 将列表转换为numpy数组 + surf_ncs = np.stack(surf_ncs_list) # [num_faces, 100, 3] + edge_ncs = edge_ncs[random_indices] # [num_faces, max_edge, 10, 3] edge_mask = edge_mask[random_indices] # [num_faces, max_edge] vertex_pos = vertex_pos[random_indices] # [num_faces, max_edge, 2, 3] diff --git a/brep2sdf/scripts/process_brep.py b/brep2sdf/scripts/process_brep.py index d70f309..77df7a2 100644 --- a/brep2sdf/scripts/process_brep.py +++ b/brep2sdf/scripts/process_brep.py @@ -246,13 +246,16 @@ def parse_solid(step_path): if points: points = np.array(points, dtype=np.float32) if len(points.shape) == 2 and points.shape[1] == 3: + # 确保每个面至少有一些点 + if len(points) < 3: # 如果点数太少，跳过这个面 + continue face_pnts.append(points) surf_bbox_wcs.append(get_bbox(shape, face)) face_explorer.Next() # Extract edge points - num_samples = 100 + num_samples = 100 # 每条边固定100个采样点 while edge_explorer.More(): edge = topods.Edge(edge_explorer.Current()) curve, first, last = BRep_Tool.Curve(edge) @@ -267,7 +270,7 @@ def parse_solid(step_path): if points: points = np.array(points, dtype=np.float32) if len(points.shape) == 2 and points.shape[1] == 3: - edge_pnts.append(points) + edge_pnts.append(points) # 这里points已经是(100, 3)形状 edge_bbox_wcs.append(get_bbox(shape, edge)) edge_explorer.Next() @@ -282,10 +285,14 @@ def parse_solid(step_path): # 获取邻接信息 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) # [num_vertices, 3] + # 转换为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 = [] @@ -298,7 +305,9 @@ def parse_solid(step_path): v1_pos, v2_pos = v2_pos, v1_pos corner_pairs.append(np.stack([v1_pos, v2_pos])) - corner_pairs = np.stack(corner_pairs) # [num_edges, 2, 3] + 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) @@ -306,18 +315,23 @@ def parse_solid(step_path): surfs_wcs, edges_wcs, surfs_ncs, edges_ncs, corner_wcs = normalize( face_pnts, edge_pnts, corner_pairs) - # Create result dictionary + # 验证归一化后的数据 + 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 + + # 创建结果字典并确保所有数组都有正确的类型 data = { - 'surf_wcs': surfs_wcs, - 'edge_wcs': edges_wcs, - 'surf_ncs': surfs_ncs, - 'edge_ncs': edges_ncs, - 'corner_wcs': corner_wcs, # [num_edges, 2, 3] - 'edgeFace_adj': edgeFace_adj, - 'edgeCorner_adj': edgeCorner_adj, - 'faceEdge_adj': faceEdge_adj, - 'surf_bbox_wcs': surf_bbox_wcs, - 'edge_bbox_wcs': edge_bbox_wcs, + '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), + 'edgeCorner_adj': edgeCorner_adj.astype(np.int32), + 'faceEdge_adj': faceEdge_adj.astype(np.int32), + '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) # 先展平再去重 } @@ -331,74 +345,43 @@ def load_step(step_path): return [reader.OneShape()] def check_data_format(data, step_file): - """检查数据格式和维度是否符合要求""" - try: - # 检查必需的键 - required_keys = [ - # 几何数据 - 'surf_wcs', 'edge_wcs', 'surf_ncs', 'edge_ncs', - 'corner_wcs', 'corner_unique', - # 拓扑关系 - 'edgeFace_adj', 'edgeCorner_adj', 'faceEdge_adj', - # 包围盒数据 - 'surf_bbox_wcs', 'edge_bbox_wcs' - ] - - # 检查键是否存在 - for key in required_keys: - if key not in data: - return False, f"Missing required key: {key}" - - # 检查几何数据 - geometry_arrays = ['surf_wcs', 'edge_wcs', 'surf_ncs', 'edge_ncs'] - for key in geometry_arrays: - if not isinstance(data[key], np.ndarray) or data[key].dtype != object: - return False, f"{key} should be a numpy array with dtype=object" - - # 检查顶点数据 - if not isinstance(data['corner_wcs'], np.ndarray) or data['corner_wcs'].dtype != np.float32: - return False, "corner_wcs should be a numpy array with dtype=float32" - if len(data['corner_wcs'].shape) != 3 or data['corner_wcs'].shape[1:] != (2, 3): - return False, f"corner_wcs should have shape (num_edges, 2, 3), got {data['corner_wcs'].shape}" - - if not isinstance(data['corner_unique'], np.ndarray) or data['corner_unique'].dtype != np.float32: - return False, "corner_unique should be a numpy array with dtype=float32" - if len(data['corner_unique'].shape) != 2 or data['corner_unique'].shape[1] != 3: - return False, f"corner_unique should have shape (N, 3), got {data['corner_unique'].shape}" - - # 检查拓扑关系 - num_faces = len(data['surf_wcs']) - num_edges = len(data['edge_wcs']) - - # 检查邻接矩阵 - adj_checks = [ - ('edgeFace_adj', (num_edges, num_faces)), - ('faceEdge_adj', (num_faces, num_edges)), - ('edgeCorner_adj', (num_edges, 2)) - ] - - for key, expected_shape in adj_checks: - if not isinstance(data[key], np.ndarray) or data[key].dtype != np.int32: - return False, f"{key} should be a numpy array with dtype=int32" - if data[key].shape != expected_shape: - return False, f"{key} shape mismatch: expected {expected_shape}, got {data[key].shape}" - - # 检查包围盒数据 - bbox_checks = [ - ('surf_bbox_wcs', (num_faces, 6)), - ('edge_bbox_wcs', (num_edges, 6)) - ] - - for key, expected_shape in bbox_checks: - if not isinstance(data[key], np.ndarray) or data[key].dtype != np.float32: - return False, f"{key} should be a numpy array with dtype=float32" - if data[key].shape != expected_shape: - return False, f"{key} shape mismatch: expected {expected_shape}, got {data[key].shape}" - - return True, "" - - except Exception as e: - return False, f"Format check failed: {str(e)}" + """检查数据格式是否正确""" + required_keys = [ + 'surf_wcs', 'edge_wcs', 'surf_ncs', 'edge_ncs', 'corner_wcs', + 'edgeFace_adj', 'edgeCorner_adj', 'faceEdge_adj', + 'surf_bbox_wcs', 'edge_bbox_wcs', 'corner_unique' + ] + + # 检查所有必需的键是否存在 + for key in required_keys: + if key not in data: + return False, f"Missing key: {key}" + + # 检查几何数据 + geometry_arrays = ['surf_wcs', 'edge_wcs', 'surf_ncs', 'edge_ncs'] + for key in geometry_arrays: + if not isinstance(data[key], np.ndarray): + return False, f"{key} should be a numpy array" + # 允许对象数组 + if data[key].dtype != object: + return False, f"{key} should be a numpy array with dtype=object" + + # 检查其他数组 + float32_arrays = ['corner_wcs', 'surf_bbox_wcs', 'edge_bbox_wcs', 'corner_unique'] + for key in float32_arrays: + if not isinstance(data[key], np.ndarray): + return False, f"{key} should be a numpy array" + if data[key].dtype != np.float32: + return False, f"{key} should be a numpy array with dtype=float32" + + int32_arrays = ['edgeFace_adj', 'edgeCorner_adj', 'faceEdge_adj'] + for key in int32_arrays: + if not isinstance(data[key], np.ndarray): + return False, f"{key} should be a numpy array" + if data[key].dtype != np.int32: + return False, f"{key} should be a numpy array with dtype=int32" + + return True, "" def process_single_step(step_path:str, output_path:str=None, timeout:int=300) -> dict: """处理单个STEP文件""" @@ -505,14 +488,13 @@ def process_furniture_step(data_path): return step_dirs - def main(): """主函数：处理多个STEP文件""" # 定义路径常量 INPUT = '/mnt/disk2/dataset/furniture/step/furniture_dataset_step/' OUTPUT = '../test_data/pkl/' RESULT = '../test_data/result/pkl/' # 用于存储成功/失败文件记录 - + # 清理输出目录 def clean_directory(directory): if os.path.exists(directory):