From 98372f171797b856adb05a7bc65833e58765c700 Mon Sep 17 00:00:00 2001 From: mckay Date: Wed, 16 Apr 2025 15:03:40 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=85=AB=E5=8F=89=E6=A0=91?= =?UTF-8?q?=E6=80=A7=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- brep2sdf/networks/encoder.py | 2 +- brep2sdf/networks/octree.py | 353 +++++++++++++++-------------------- brep2sdf/train.py | 2 +- 3 files changed, 157 insertions(+), 200 deletions(-) diff --git a/brep2sdf/networks/encoder.py b/brep2sdf/networks/encoder.py index 8c3d4d1..9dfb0e8 100644 --- a/brep2sdf/networks/encoder.py +++ b/brep2sdf/networks/encoder.py @@ -88,7 +88,7 @@ class Encoder(nn.Module): while stack: node, path = stack.pop() - if node.is_leaf(): + if node._is_leaf: # 如果是叶子节点，初始化参数 param_name = f"leaf_{path}" self._leaf_parameters.append(nn.Parameter(torch.randn(8, self.feature_dim))) # 8个顶点的特征 diff --git a/brep2sdf/networks/octree.py b/brep2sdf/networks/octree.py index 4b2a308..2d3929f 100644 --- a/brep2sdf/networks/octree.py +++ b/brep2sdf/networks/octree.py @@ -1,5 +1,3 @@ - - from typing import Tuple, List, cast, Dict, Any, Tuple import torch @@ -9,7 +7,7 @@ import numpy as np from brep2sdf.utils.logger import logger -def bbox_intersect(bbox1: torch.Tensor, bbox2: torch.Tensor) -> bool: +def bbox_intersect(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor: """判断两个轴对齐包围盒(AABB)是否相交 参数: @@ -17,7 +15,7 @@ def bbox_intersect(bbox1: torch.Tensor, bbox2: torch.Tensor) -> bool: bbox2: 同bbox1格式 返回: - bool: 两包围盒是否相交(包括刚好接触的情况) + torch.Tensor: 两包围盒是否相交(包括刚好接触的情况) """ assert bbox1.shape == (6,) and bbox2.shape == (6,), "输入必须是形状为(6,)的张量" @@ -29,191 +27,142 @@ def bbox_intersect(bbox1: torch.Tensor, bbox2: torch.Tensor) -> bool: return torch.all((max1 >= min2) & (max2 >= min1)) class OctreeNode(nn.Module): - device=None - surf_bbox = None - def __init__(self, bbox: torch.Tensor,face_indices: np.ndarray, max_depth: int = 5, surf_bbox:torch.Tensor = None): + def __init__(self, bbox: torch.Tensor, face_indices: np.ndarray, max_depth: int = 5, surf_bbox: torch.Tensor = None): super().__init__() - self.bbox = bbox # 节点的边界框 - self.max_depth: int = max_depth # 最大深度，当这个为0时，表示已经到达最大深度，不可再分子节点 - self.child_nodes: torch.nn.ModuleList = torch.nn.ModuleList() # 子节点列表 - self.face_indices = face_indices - self.param_key = "" - #self.patch_feature_volume = None # 补丁特征体积，only leaf has + # 静态张量存储节点信息 + self.register_buffer('bbox', bbox) # 当前节点的边界框 + self.register_buffer('node_bboxes', None) # 所有节点的边界框 + self.register_buffer('parent_indices', None) # 父节点索引 + self.register_buffer('child_indices', None) # 子节点索引 + self.register_buffer('is_leaf_mask', None) # 叶子节点标记 + self.register_buffer('face_indices', torch.from_numpy(face_indices).to(bbox.device)) # 面片索引张量 + self.register_buffer('surf_bbox', surf_bbox) # 面片边界框 + + self.max_depth = max_depth + self.param_key = "" self._is_leaf = True - #print(f"box shape: {self.bbox.shape}") - if surf_bbox is not None: - if not isinstance(surf_bbox, torch.Tensor): - raise TypeError( - f"surf_bbox 必须是 torch.Tensor 类型，但得到 {type(surf_bbox)}" - ) - if surf_bbox.dim() != 2 or surf_bbox.shape[1] != 6: - raise ValueError( - f"surf_bbox 应为二维张量且形状为 (N,6)，但得到 {surf_bbox.shape}" - ) - OctreeNode.surf_bbox = surf_bbox # NOTE: 只在根节点时创建 - OctreeNode.device = bbox.device - - def is_leaf(self): - # Check if self.child——nodes is None before calling len() - return self._is_leaf - - def set_param_key(self, k): + @torch.jit.export + def set_param_key(self, k: str) -> None: self.param_key = k - - def conduct_tree(self): - if self.max_depth <= 0 or self.face_indices.shape[0] <= 2: - # 达到最大深度 or 一个单元格至多只有两个面 - return - self.subdivide() - - - def subdivide(self): + + @torch.jit.export + def build_static_tree(self) -> None: + """构建静态八叉树结构""" + # 预计算所有可能的节点数量，确保结果为整数 + total_nodes = int(sum(8**i for i in range(self.max_depth + 1))) - #min_x, min_y, min_z, max_x, max_y, max_z = self.bbox - # 使用索引操作替代解包 - min_coords = self.bbox[:3] # [min_x, min_y, min_z] - max_coords = self.bbox[3:] # [max_x, max_y, max_z] - - # 计算中间点 - mid_coords = (min_coords + max_coords) / 2 + # 初始化静态张量，使用整数列表作为形状参数 + self.node_bboxes = torch.zeros([int(total_nodes), 6], device=self.bbox.device) + self.parent_indices = torch.full([int(total_nodes)], -1, dtype=torch.long, device=self.bbox.device) + self.child_indices = torch.full([int(total_nodes), 8], -1, dtype=torch.long, device=self.bbox.device) + self.is_leaf_mask = torch.zeros([int(total_nodes)], dtype=torch.bool, device=self.bbox.device) - # 提取 min_x, min_y, min_z, mid_x, mid_y, mid_z - min_x, min_y, min_z = min_coords[0], min_coords[1], min_coords[2] - mid_x, mid_y, mid_z = mid_coords[0], mid_coords[1], mid_coords[2] - max_x, max_y, max_z = max_coords[0], max_coords[1], max_coords[2] - - # 生成 8 个子包围盒 - child_bboxes = torch.stack([ - torch.cat([min_coords, mid_coords]), # 前下左 - torch.cat([torch.tensor([mid_x, min_y, min_z], device=self.bbox.device), - torch.tensor([max_x, mid_y, mid_z], device=self.bbox.device)]), # 前下右 - torch.cat([torch.tensor([min_x, mid_y, min_z], device=self.bbox.device), - torch.tensor([mid_x, max_y, mid_z], device=self.bbox.device)]), # 前上左 - torch.cat([torch.tensor([mid_x, mid_y, min_z], device=self.bbox.device), - torch.tensor([max_x, max_y, mid_z], device=self.bbox.device)]), # 前上右 - torch.cat([torch.tensor([min_x, min_y, mid_z], device=self.bbox.device), - torch.tensor([mid_x, mid_y, max_z], device=self.bbox.device)]), # 后下左 - torch.cat([torch.tensor([mid_x, min_y, mid_z], device=self.bbox.device), - torch.tensor([max_x, mid_y, max_z], device=self.bbox.device)]), # 后下右 - torch.cat([torch.tensor([min_x, mid_y, mid_z], device=self.bbox.device), - torch.tensor([mid_x, max_y, max_z], device=self.bbox.device)]), # 后上左 - torch.cat([torch.tensor([mid_x, mid_y, mid_z], device=self.bbox.device), - torch.tensor([max_x, max_y, max_z], device=self.bbox.device)]) # 后上右 - ]) + # 使用队列进行广度优先遍历 + queue = [(0, self.bbox, self.face_indices)] # (node_idx, bbox, face_indices) + current_idx = 0 - # 为每个子包围盒创建子节点，并分配相交的面 - for bbox in child_bboxes: - # 找到与子包围盒相交的面 - intersecting_faces = [] - for face_idx in self.face_indices: - face_bbox = OctreeNode.surf_bbox[face_idx] - if bbox_intersect(bbox, face_bbox): - intersecting_faces.append(face_idx) - #print(f"{bbox}: {intersecting_faces}") + while queue: + node_idx, bbox, faces = queue.pop(0) + self.node_bboxes[node_idx] = bbox + + if faces.shape[0] <= 2 or current_idx >= self.max_depth: + self.is_leaf_mask[node_idx] = True + continue + + # 计算子节点边界框 + min_coords = bbox[:3] + max_coords = bbox[3:] + mid_coords = (min_coords + max_coords) / 2 - child_node = OctreeNode( - bbox=bbox, - face_indices=np.array(intersecting_faces), - max_depth=self.max_depth - 1 - ) - child_node.conduct_tree() - self.child_nodes.append(child_node) + # 生成8个子节点 + child_bboxes = self._generate_child_bboxes(min_coords, mid_coords, max_coords) + + # 为每个子节点分配面片 + for i, child_bbox in enumerate(child_bboxes): + child_idx = current_idx + 1 + current_idx += 1 + + # 找到与子包围盒相交的面 + intersecting_faces = [] + for face_idx in faces: + face_bbox = self.surf_bbox[face_idx] + if bbox_intersect(child_bbox, face_bbox).item(): + intersecting_faces.append(face_idx) + + # 更新节点关系 + self.parent_indices[child_idx] = node_idx + self.child_indices[node_idx, i] = child_idx + + # 将子节点加入队列 + if intersecting_faces: + queue.append((child_idx, child_bbox, torch.tensor(intersecting_faces, device=self.bbox.device))) + + @torch.jit.export + def _generate_child_bboxes(self, min_coords: torch.Tensor, mid_coords: torch.Tensor, max_coords: torch.Tensor) -> torch.Tensor: + """生成8个子节点的边界框""" + child_bboxes = torch.zeros([8, 6], device=self.bbox.device) - self._is_leaf = False - - def get_child_index(self, query_point: torch.Tensor) -> int: - """ - 计算点所在子节点的索引 - :param query_point: 待检查的点，格式为 (x, y, z) - :return: 子节点的索引，范围从 0 到 7 - """ - # 确保 query_point 和 bbox 在同一设备上 - query_point = query_point.to(self.bbox.device) - - # 提取 bbox 的最小和最大坐标 - min_coords = self.bbox[:3] # [min_x, min_y, min_z] - max_coords = self.bbox[3:] # [max_x, max_y, max_z] - - # 计算中间点 - mid_coords = (min_coords + max_coords) / 2 - - # 使用布尔比较结果计算索引 - index = ((query_point >= mid_coords) << torch.arange(3, device=self.bbox.device)).sum() - - return index.item() - - def find_leaf(self, query_point: torch.Tensor) -> Tuple[torch.Tensor, str, bool]: + # 使用向量化操作生成所有子节点边界框 + child_bboxes[0] = torch.cat([min_coords, mid_coords]) # 前下左 + child_bboxes[1] = torch.cat([torch.stack([mid_coords[0], min_coords[1], min_coords[2]]), + torch.stack([max_coords[0], mid_coords[1], mid_coords[2]])]) # 前下右 + child_bboxes[2] = torch.cat([torch.stack([min_coords[0], mid_coords[1], min_coords[2]]), + torch.stack([mid_coords[0], max_coords[1], mid_coords[2]])]) # 前上左 + child_bboxes[3] = torch.cat([torch.stack([mid_coords[0], mid_coords[1], min_coords[2]]), + torch.stack([max_coords[0], max_coords[1], mid_coords[2]])]) # 前上右 + child_bboxes[4] = torch.cat([torch.stack([min_coords[0], min_coords[1], mid_coords[2]]), + torch.stack([mid_coords[0], mid_coords[1], max_coords[2]])]) # 后下左 + child_bboxes[5] = torch.cat([torch.stack([mid_coords[0], min_coords[1], mid_coords[2]]), + torch.stack([max_coords[0], mid_coords[1], max_coords[2]])]) # 后下右 + child_bboxes[6] = torch.cat([torch.stack([min_coords[0], mid_coords[1], mid_coords[2]]), + torch.stack([mid_coords[0], max_coords[1], max_coords[2]])]) # 后上左 + child_bboxes[7] = torch.cat([mid_coords, max_coords]) # 后上右 + + return child_bboxes + + @torch.jit.export + def find_leaf(self, query_points: torch.Tensor) -> Tuple[torch.Tensor, str, bool]: """ - 查找包含给定点的叶子节点，并返回其信息（以元组形式） - :param query_point: 待查找的点 + 查找包含给定点的叶子节点，并返回其信息 + :param query_points: 待查找的点，形状为 (3,) :return: 包含叶子节点信息的元组 (bbox, param_key, is_leaf) """ - # 如果当前节点是叶子节点，返回其信息 - if self._is_leaf: - #logger.info(f"{self.bbox}, {self.param_key}, {True}") - return (self.bbox, self.param_key, True) - - # 计算查询点所在的子节点索引 - index = self.get_child_index(query_point) - - # 遍历子节点列表，找到对应的子节点 - for i, child_node in enumerate(self.child_nodes): - if i == index and child_node is not None: - # 递归调用子节点的 find_leaf 方法 - result = child_node.find_leaf(query_point) + # 确保输入是单个点 + if query_points.dim() != 1 or query_points.shape[0] != 3: + raise ValueError(f"query_points 必须是形状为 (3,) 的张量，但得到 {query_points.shape}") + + current_idx = torch.tensor(0, dtype=torch.long, device=query_points.device) + max_iterations = 1000 # 防止无限循环 + iteration = 0 + + while iteration < max_iterations: + # 获取当前节点的叶子状态 + if self.is_leaf_mask[current_idx].item(): + return self.node_bboxes[current_idx], self.param_key, True + + # 计算子节点索引 + child_idx = self._get_child_indices(query_points.unsqueeze(0), + self.node_bboxes[current_idx].unsqueeze(0)) + + # 获取下一个要访问的节点 + next_idx = self.child_indices[current_idx, child_idx[0]] + + # 检查索引是否有效 + if next_idx == -1: + raise IndexError(f"Invalid child node index: {child_idx[0]}") - # 确保返回值是一个元组 - assert isinstance(result, tuple), f"Unexpected return type: {type(result)}" - return result - - # 如果找不到有效的子节点，抛出异常 - raise IndexError(f"Invalid child node index: {index}") - ''' - try: - # 直接访问子节点，不进行显式检查 - return self.child_nodes[index].find_leaf(query_point) - except IndexError as e: - # 记录错误日志并重新抛出异常 - logger.error( - f"Error accessing child node: {e}. " - f"Query point: {query_point.cpu().numpy().tolist()}, " - f"Node bbox: {self.bbox.cpu().numpy().tolist()}, " - f"Depth info: {self.max_depth}" - ) - raise e - ''' - - ''' - def get_feature_vector(self, query_point:torch.Tensor): - """ - 预测给定点的 SDF 值 - :param point: 待预测的点，格式为 (x, y, z) - :return: 预测的 SDF 值 - """ - # 将点转换为 numpy 数组 - - # 从根节点开始递归查找包含该点的叶子节点 - if self._is_leaf: - return self.trilinear_interpolation(query_point) - else: - index = self.get_child_index(query_point) - try: - # 直接访问子节点，不进行显式检查 - return self.child_nodes[index].get_feature_vector(query_point) - except IndexError as e: - # 记录错误日志并重新抛出异常 - logger.error( - f"Error accessing child node: {e}. " - f"Query point: {query_point.cpu().numpy().tolist()}, " - f"Node bbox: {self.bbox.cpu().numpy().tolist()}, " - f"Depth info: {self.max_depth}" - ) - raise e - ''' - - + current_idx = next_idx + iteration += 1 + + # 如果达到最大迭代次数，返回当前节点的信息 + return self.node_bboxes[current_idx], self.param_key, bool(self.is_leaf_mask[current_idx].item()) + @torch.jit.export + def _get_child_indices(self, points: torch.Tensor, bboxes: torch.Tensor) -> torch.Tensor: + """批量计算点所在的子节点索引""" + mid_points = (bboxes[:, :3] + bboxes[:, 3:]) / 2 + return ((points >= mid_points) << torch.arange(3, device=points.device)).sum(dim=1) def print_tree(self, depth: int = 0, max_print_depth: int = None) -> None: """ @@ -233,33 +182,41 @@ class OctreeNode(nn.Module): # 打印面片信息（如果有） if self.face_indices is not None: - print(f"{indent} Face indices: {self.face_indices.tolist()}") - print(f"{indent} len child_nodes: {len(self.child_nodes)}") + print(f"{indent} Face indices: {self.face_indices.cpu().numpy().tolist()}") + print(f"{indent} Child indices: {self.child_indices.cpu().numpy().tolist()}") - # 递归打印子节点 - for i, child in enumerate(self.child_nodes): - print(f"{indent} Child {i}:") - child.print_tree(depth + 1, max_print_depth) + # 打印子节点信息 + if self.child_indices is not None: + for i in range(8): + child_idx = self.child_indices[0, i].item() + if child_idx != -1: + print(f"{indent} Child {i}: Node {child_idx}") def __getstate__(self): """支持pickle序列化""" - return self._serialize_node(self) + state = { + 'bbox': self.bbox, + 'node_bboxes': self.node_bboxes, + 'parent_indices': self.parent_indices, + 'child_indices': self.child_indices, + 'is_leaf_mask': self.is_leaf_mask, + 'face_indices': self.face_indices, + 'surf_bbox': self.surf_bbox, + 'max_depth': self.max_depth, + 'param_key': self.param_key, + '_is_leaf': self._is_leaf + } + return state def __setstate__(self, state): """支持pickle反序列化""" - self = self._deserialize_node(state) - - def _serialize_node(self, node): - return { - 'bbox': node.bbox, - 'is_leaf': node._is_leaf, - 'child_nodes': [self._serialize_node(c) for c in node.child_nodes], - 'param_key': node.param_key - } - - def _deserialize_node(self, data): - node = OctreeNode(data['bbox'], 0) # max_depth会在encoder中重建 - node._is_leaf = data['is_leaf'] - node.param_key = data['param_key'] - node.child_nodes = [self._deserialize_node(c) for c in data['child_nodes']] - return node \ No newline at end of file + self.bbox = state['bbox'] + self.node_bboxes = state['node_bboxes'] + self.parent_indices = state['parent_indices'] + self.child_indices = state['child_indices'] + self.is_leaf_mask = state['is_leaf_mask'] + self.face_indices = state['face_indices'] + self.surf_bbox = state['surf_bbox'] + self.max_depth = state['max_depth'] + self.param_key = state['param_key'] + self._is_leaf = state['_is_leaf'] \ No newline at end of file diff --git a/brep2sdf/train.py b/brep2sdf/train.py index 474f856..e589753 100644 --- a/brep2sdf/train.py +++ b/brep2sdf/train.py @@ -140,7 +140,7 @@ class Trainer: ) #print(surf_bbox) logger.info("starting octree conduction") - self.root.conduct_tree() + self.root.build_static_tree() logger.info("complete octree conduction") #self.root.print_tree(0)