From 85a581be1771d3623b1456f2ac7ec8d4ea640793 Mon Sep 17 00:00:00 2001
From: mckay <wchpub@163.com>
Date: Mon, 18 Nov 2024 23:05:33 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E6=94=B9=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=8A=A0=E8=BD=BDprocess=5Fbrep=5Fdata=EF=BC=8C=E9=98=B2?=
 =?UTF-8?q?=E6=AD=A2=E7=BB=B4=E5=BA=A6=E8=B6=85=E8=BF=87=E6=9C=80=E5=A4=A7?=
 =?UTF-8?q?=EF=BC=88=E7=94=A8=E6=88=AA=E5=8F=96=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 brep2sdf/data/utils.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/brep2sdf/data/utils.py b/brep2sdf/data/utils.py
index 159e289..354aa83 100644
--- a/brep2sdf/data/utils.py
+++ b/brep2sdf/data/utils.py
@@ -1169,15 +1169,19 @@ def process_brep_data(data: dict, max_face: int, max_edge: int, bbox_scaled: flo
     vertex_pos = vertex_pos[random_indices]  # [num_faces, max_edge, 2, 3]
     
     # 填充到最大面数
-    surf_pos = pad_zero(surf_pos, max_face)     # [max_face, 6]
+    surf_pos, face_mask = pad_zero(surf_pos, max_face, return_mask=True)  # [max_face, 6]
     surf_ncs = pad_zero(surf_ncs, max_face)     # [max_face, 100, 3]
     edge_pos = pad_zero(edge_pos, max_face)     # [max_face, max_edge, 6]
     edge_ncs = pad_zero(edge_ncs, max_face)     # [max_face, max_edge, 10, 3]
     vertex_pos = pad_zero(vertex_pos, max_face)  # [max_face, max_edge, 2, 3]
     
-    # 扩展边掩码
-    padding = np.zeros((max_face-len(edge_mask), *edge_mask.shape[1:])) == 0  # [max_face-num_faces, max_edge]
-    edge_mask = np.concatenate([edge_mask, padding], 0)  # [max_face, max_edge]
+    # 扩展边掩码 - 使用face_mask来创建新的edge_mask
+    if len(edge_mask) > max_face:
+        edge_mask = edge_mask[:max_face]
+    else:
+        # 创建填充掩码
+        padding = np.zeros((max_face-len(edge_mask), max_edge), dtype=bool)
+        edge_mask = np.concatenate([edge_mask, padding], axis=0)  # [max_face, max_edge]
     
     # 转换为张量并返回
     if data_class is not None: