From 94897c10b68b8b9dc1648771524098eac1802e54 Mon Sep 17 00:00:00 2001 From: mckay Date: Fri, 15 Nov 2024 00:12:52 +0800 Subject: [PATCH] data laod --- README.md | 49 ++++++++- brep2sdf/data/data.py | 181 +++++++++++++++++++++++++------ brep2sdf/scripts/process_brep.py | 31 +++++- brep2sdf/scripts/test_pkl.py | 0 4 files changed, 224 insertions(+), 37 deletions(-) create mode 100644 brep2sdf/scripts/test_pkl.py diff --git a/README.md b/README.md index c8a9427..6726219 100644 --- a/README.md +++ b/README.md @@ -1 +1,48 @@ -# brep2sdf \ No newline at end of file +## 做什么事情 +学习 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里面 面点的处理方式 diff --git a/brep2sdf/data/data.py b/brep2sdf/data/data.py index eec7244..51ccbfd 100644 --- a/brep2sdf/data/data.py +++ b/brep2sdf/data/data.py @@ -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: diff --git a/brep2sdf/scripts/process_brep.py b/brep2sdf/scripts/process_brep.py index 244488f..6e557e0 100644 --- a/brep2sdf/scripts/process_brep.py +++ b/brep2sdf/scripts/process_brep.py @@ -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") diff --git a/brep2sdf/scripts/test_pkl.py b/brep2sdf/scripts/test_pkl.py new file mode 100644 index 0000000..e69de29