From 2b5860a7e7b3f4436bbe0b13e39620972618e83f Mon Sep 17 00:00:00 2001
From: mckay <wchpub@163.com>
Date: Fri, 22 Nov 2024 00:40:23 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0sdf=20dataloader=E9=9A=8F?=
=?UTF-8?q?=E6=9C=BA=E9=87=87=E6=A0=B7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
brep2sdf/data/data.py | 30 +++++++++++++++++++++++++++++-
1 file changed, 29 insertions(+), 1 deletion(-)
diff --git a/brep2sdf/data/data.py b/brep2sdf/data/data.py
index 0e1b322..a893de3 100644
--- a/brep2sdf/data/data.py
+++ b/brep2sdf/data/data.py
@@ -271,7 +271,7 @@ class BRepSDFDataset(Dataset):
raise
def _load_sdf_file(self, sdf_path):
- """加载和处理SDF数据"""
+ """加载和处理SDF数据，并进行随机采样"""
try:
# 加载SDF值
sdf_data = np.load(sdf_path)
@@ -285,7 +285,35 @@ class BRepSDFDataset(Dataset):
if sdf_pos.shape[1] != 4 or sdf_neg.shape[1] != 4:
raise ValueError(f"Invalid SDF data shape: pos={sdf_pos.shape}, neg={sdf_neg.shape}")
+ # 随机采样
+ max_points = self.config.data.num_query_points # 例如4096
+
+ # 确保正负样本均衡
+ num_pos = min(max_points // 2, sdf_pos.shape[0])
+ num_neg = min(max_points // 2, sdf_neg.shape[0])
+
+ # 随机采样正样本
+ if sdf_pos.shape[0] > num_pos:
+ pos_indices = np.random.choice(sdf_pos.shape[0], num_pos, replace=False)
+ sdf_pos = sdf_pos[pos_indices]
+
+ # 随机采样负样本
+ if sdf_neg.shape[0] > num_neg:
+ neg_indices = np.random.choice(sdf_neg.shape[0], num_neg, replace=False)
+ sdf_neg = sdf_neg[neg_indices]
+
+ # 合并数据
sdf_np = np.concatenate([sdf_pos, sdf_neg], axis=0)
+
+ # 再次随机打乱
+ np.random.shuffle(sdf_np)
+
+ # 如果总点数仍然超过最大限制，再次采样
+ if sdf_np.shape[0] > max_points:
+ indices = np.random.choice(sdf_np.shape[0], max_points, replace=False)
+ sdf_np = sdf_np[indices]
+
+ logger.debug(f"Sampled SDF points: {sdf_np.shape[0]} (max: {max_points})")
return torch.from_numpy(sdf_np.astype(np.float32))
except Exception as e: