wch
/
brep2sdf
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

data laod

main
mckay 4 months ago
parent
e5d9855941
commit
94897c10b6
4 changed files with 224 additions and 37 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 49
      README.md
  2. 181
      brep2sdf/data/data.py
  3. 31
      brep2sdf/scripts/process_brep.py
  4. 0
      brep2sdf/scripts/test_pkl.py

49
README.md
View File

@ -1 +1,48 @@
# brep2sdf
## 做什么事情
学习 brep 的sdf表示
## 意义
- brep方便控制
- sdf方便运算和智能生成
## 怎么做
### 数据准备
数据源： furniture
文件结构
brep( .step )
train
chair
chair_1.brep
chair_2.brep
test
val
sdf ( .npz )
train
chair
chair_1.sdf.npz
chair_2.sdf.npz
数据结构
### brep to mesh
基本思路：采样，然后三角化
- 重采样
《Mesh Generation from CAD Models》: sharp edge
### mesh to sdf
基本思路：体素，然后积分
docker
```
docker run -it -v /mnt/disk2/dataset:/data -v $(pwd):/app --name brep_sdf ubuntu:22.04
```
# 一些点
dataset里面 面点的处理方式

181
brep2sdf/data/data.py
View File

@ -29,49 +29,161 @@ class BRepSDFDataset(Dataset):
raise ValueError(f"SDF directory not found: {self.sdf_dir}")
# 加载数据列表
self.data_list = self._load_data_list()
self.brep_data_list = self._load_data_list(self.brep_dir)
self.sdf_data_list = self._load_data_list(self.sdf_dir)
# 检查数据集是否为空
if len(self.data_list) == 0:
raise ValueError(f"No valid data found in {split} set")
if len(self.brep_data_list) == 0 :
raise ValueError(f"No valid brep data found in {split} set")
if len(self.sdf_data_list) == 0:
raise ValueError(f"No valid sdf data found in {split} set")
logger.info(f"Loaded {split} dataset with {len(self.data_list)} samples")
logger.info(f"Loaded {split} dataset with {len(self.brep_data_list)} samples")
def _load_data_list(self):
# data_dir 为 self.brep_dir or sdf_dir
def _load_data_list(self, data_dir):
data_list = []
for sample_dir in os.listdir(self.brep_dir):
sample_path = os.path.join(self.brep_dir, sample_dir)
if os.path.isdir(sample_path):
data_list.append(sample_path)
for sample_file in os.listdir(data_dir):
path = os.path.join(data_dir, sample_file)
data_list.append(path)
#logger.info(data_list)
return data_list
def __len__(self):
return len(self.data_list)
return len(self.brep_data_list)
def __getitem__(self, idx):
sample_path = self.data_list[idx]
# 解析 .step 文件
step_file = os.path.join(sample_path, 'model.step')
brep_features = self._parse_step_file(step_file)
# 加载 .sdf 文件
sdf_file = os.path.join(sample_path, 'sdf.npy')
sdf = np.load(sdf_file)
# 转换为 torch 张量
brep_features = torch.tensor(brep_features, dtype=torch.float32)
sdf = torch.tensor(sdf, dtype=torch.float32)
"""获取单个数据样本"""
brep_path = self.brep_data_list[idx]
sdf_path = self.sdf_data_list[idx]
return {
'brep_features': brep_features,
'sdf': sdf
}
try:
# 获取文件名(不含扩展名)作为sample name
name = os.path.splitext(os.path.basename(brep_path))[0]
# 加载B-rep和SDF数据
brep_data = self._load_brep_file(brep_path)
sdf_data = self._load_sdf_file(sdf_path)
# 修改返回格式，将sdf_data作为一个键值对添加
return {
'name': name,
**brep_data, # 解包B-rep特征
'sdf': sdf_data # 添加SDF数据作为一个键
}
except Exception as e:
logger.error(f"Error loading sample from {brep_path}: {str(e)}")
logger.error("Data structure:")
if 'brep_data' in locals():
for key, value in brep_data.items():
if isinstance(value, np.ndarray):
logger.error(f" {key}: type={type(value)}, dtype={value.dtype}, shape={value.shape}")
raise
def _parse_step_file(self, step_file):
# 解析 .step 文件的逻辑
# 返回 B-rep 特征
pass
def _load_brep_file(self, brep_path):
"""加载B-rep特征文件"""
try:
with open(brep_path, 'rb') as f:
brep_data = pickle.load(f)
features = {}
# 1. 处理几何数据（不等长序列）
for key in ['surf_wcs', 'surf_ncs', 'edge_wcs', 'edge_ncs']:
if key in brep_data:
try:
features[key] = [
torch.from_numpy(np.array(x, dtype=np.float32))
for x in brep_data[key]
]
except Exception as e:
logger.error(f"Error converting {key}:")
logger.error(f" Type: {type(brep_data[key])}")
if isinstance(brep_data[key], list):
logger.error(f" List length: {len(brep_data[key])}")
if len(brep_data[key]) > 0:
logger.error(f" First element type: {type(brep_data[key][0])}")
logger.error(f" First element shape: {brep_data[key][0].shape}")
logger.error(f" First element dtype: {brep_data[key][0].dtype}")
raise ValueError(f"Failed to convert {key}: {str(e)}")
# 2. 处理固定形状的数据
for key in ['corner_wcs', 'corner_unique', 'surf_bbox_wcs', 'edge_bbox_wcs']:
if key in brep_data:
try:
data = np.array(brep_data[key], dtype=np.float32)
features[key] = torch.from_numpy(data)
except Exception as e:
logger.error(f"Error converting {key}:")
logger.error(f" Type: {type(brep_data[key])}")
if isinstance(brep_data[key], np.ndarray):
logger.error(f" Shape: {brep_data[key].shape}")
logger.error(f" dtype: {brep_data[key].dtype}")
raise ValueError(f"Failed to convert {key}: {str(e)}")
# 3. 处理邻接矩阵
for key in ['edgeFace_adj', 'edgeCorner_adj', 'faceEdge_adj']:
if key in brep_data:
try:
data = np.array(brep_data[key], dtype=np.int32)
features[key] = torch.from_numpy(data)
except Exception as e:
logger.error(f"Error converting {key}:")
logger.error(f" Type: {type(brep_data[key])}")
if isinstance(brep_data[key], np.ndarray):
logger.error(f" Shape: {brep_data[key].shape}")
logger.error(f" dtype: {brep_data[key].dtype}")
raise ValueError(f"Failed to convert {key}: {str(e)}")
return features
except Exception as e:
logger.error(f"\nError loading B-rep file: {brep_path}")
logger.error(f"Error message: {str(e)}")
# 打印完整的数据结构信息
if 'brep_data' in locals():
logger.error("\nComplete data structure:")
for key, value in brep_data.items():
logger.error(f"\n{key}:")
logger.error(f" Type: {type(value)}")
if isinstance(value, np.ndarray):
logger.error(f" Shape: {value.shape}")
logger.error(f" dtype: {value.dtype}")
elif isinstance(value, list):
logger.error(f" List length: {len(value)}")
if len(value) > 0:
logger.error(f" First element type: {type(value[0])}")
if isinstance(value[0], np.ndarray):
logger.error(f" First element shape: {value[0].shape}")
logger.error(f" First element dtype: {value[0].dtype}")
raise
def _load_sdf_file(self, sdf_path):
"""加载和处理SDF数据"""
try:
# 加载SDF值
sdf_data = np.load(sdf_path)
if 'pos' not in sdf_data or 'neg' not in sdf_data:
raise ValueError("Missing pos/neg data in SDF file")
sdf_pos = sdf_data['pos'] # (N1, 4)
sdf_neg = sdf_data['neg'] # (N2, 4)
# 添加数据验证
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}")
sdf_np = np.concatenate([sdf_pos, sdf_neg], axis=0)
return torch.from_numpy(sdf_np.astype(np.float32))
except Exception as e:
logger.error(f"Error loading SDF from {sdf_path}")
logger.error(f"Error type: {type(e).__name__}")
logger.error(f"Error message: {str(e)}")
raise
def test_dataset():
"""测试数据集功能"""
@ -88,11 +200,10 @@ def test_dataset():
logger.info(f"Split: {split}")
# ... (其余测试代码保持不变) ...
dataset = BRepSDFDataset(brep_data_dir='/home/wch/myDeepSDF/test_data/pkl', sdf_data_dir='/home/wch/myDeepSDF/test_data/sdf', split='train')
dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True)
dataset = BRepSDFDataset(brep_dir='/home/wch/brep2sdf/test_data/pkl', sdf_dir='/home/wch/brep2sdf/test_data/sdf', split='train')
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
for batch in dataloader:
print(batch['brep_features'].shape)
print(batch['sdf'].shape)
break
except Exception as e:

31
brep2sdf/scripts/process_brep.py
View File

@ -202,8 +202,36 @@ def get_bbox(shape, subshape):
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."""
"""
返回值：一个dict，包括
几何关系
# 面片数据
'surf_wcs': surfs_wcs, # list of np.array(N, 3), 每个面片的点云坐标
'surf_ncs': surfs_ncs, # list of np.array(N, 3), 每个面片的法向量
# 边数据 (num_samples=100)
'edge_wcs': edges_wcs, # list of np.array(100, 3), 每条边的采样点坐标
'edge_ncs': edges_ncs, # list of np.array(100, 3), 每条边的法向量
# 顶点数据
'corner_wcs': corner_wcs.astype(np.float32), # np.array(N, 3), 所有顶点坐标
'corner_unique': np.unique(corner_wcs, axis=0).astype(np.float32), # np.array(M, 3), 去重后的顶点坐标
拓扑关系
# 邻接矩阵，都是int32类型
'edgeFace_adj': edgeFace_adj, # np.array(num_edges, num_faces), 边-面邻接关系
'edgeCorner_adj': edgeCorner_adj, # np.array(num_edges, 2), 边-顶点邻接关系
'faceEdge_adj': faceEdge_adj, # np.array(num_faces, num_edges), 面-边邻接关系
包围盒数据
# 包围盒坐标，float32类型
'surf_bbox_wcs': surf_bbox_wcs, # np.array(num_faces, 6), 每个面的包围盒 [xmin,ymin,zmin,xmax,ymax,zmax]
'edge_bbox_wcs': edge_bbox_wcs, # np.array(num_edges, 6), 每条边的包围盒 [xmin,ymin,zmin,xmax,ymax,zmax]
"""
# Load STEP file
reader = STEPControl_Reader()
status = reader.ReadFile(step_path)
@ -488,4 +516,5 @@ def main():
logger.warning("No files were processed")
if __name__ == '__main__':
main()
#main()
test("/mnt/disk2/dataset/furniture/step/furniture_dataset_step/train/bathtub_0004.step", "/home/wch/brep2sdf/test_data/pkl/train/bathtub_0004.pkl")

0
brep2sdf/scripts/test_pkl.py
View File

Write Preview
Loading…
Cancel
Save