Browse Source

fix: 修改了check_data_format函数,正确检查对象数组类型

在parse_solid函数中确保所有数组都有正确的类型
对点云数据使用对象数组类型
对其他数据使用正确的数值类型(float32或int32)
final
mckay 7 months ago
parent
commit
295e707e2b
  1. 23
      brep2sdf/data/utils.py
  2. 124
      brep2sdf/scripts/process_brep.py

23
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]

124
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,75 +345,44 @@ 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',
# 拓扑关系
'surf_wcs', 'edge_wcs', 'surf_ncs', 'edge_ncs', 'corner_wcs',
'edgeFace_adj', 'edgeCorner_adj', 'faceEdge_adj',
# 包围盒数据
'surf_bbox_wcs', 'edge_bbox_wcs'
'surf_bbox_wcs', 'edge_bbox_wcs', 'corner_unique'
]
# 检查键是否存在
# 检查所有必需的键是否存在
for key in required_keys:
if key not in data:
return False, f"Missing required key: {key}"
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) or data[key].dtype != object:
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"
# 检查顶点数据
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))
]
# 检查其他数组
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"
for key, expected_shape in adj_checks:
if not isinstance(data[key], np.ndarray) or data[key].dtype != np.int32:
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"
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)}"
def process_single_step(step_path:str, output_path:str=None, timeout:int=300) -> dict:
"""处理单个STEP文件"""
try:
@ -505,7 +488,6 @@ def process_furniture_step(data_path):
return step_dirs
def main():
"""主函数:处理多个STEP文件"""
# 定义路径常量

Loading…
Cancel
Save