feat: update algorithm integrations and point cloud overlays

This commit is contained in:
2026-06-30 16:46:36 +08:00
parent ddca7e5436
commit 7e2e56a50d
64 changed files with 4356 additions and 2781 deletions

View File

@@ -0,0 +1,17 @@
"""空间索引模块 - 提供多种ROI提取方法"""
from .base import SpatialIndex
from .linear_scan import LinearScanIndex
from .grid_index import GridIndex
from .grid_index_optimized import GridIndexOptimized
from .octree_index import OctreeIndex
from .kdtree_index import KDTreeIndex
__all__ = [
'SpatialIndex',
'LinearScanIndex',
'GridIndex',
'GridIndexOptimized',
'OctreeIndex',
'KDTreeIndex',
]

View File

@@ -0,0 +1,39 @@
"""空间索引基类"""
from abc import ABC, abstractmethod
import numpy as np
class SpatialIndex(ABC):
"""空间索引抽象基类"""
def __init__(self, points: np.ndarray):
"""
初始化索引
Args:
points: 点云数组 (N, 3)
"""
self.points = points
self.build_time = 0.0 # 构建时间(秒)
@abstractmethod
def query_bbox(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
"""
查询包围盒内的点
Args:
bbox_min: [x_min, y_min, z_min]
bbox_max: [x_max, y_max, z_max]
Returns:
包围盒内的点的索引数组
"""
pass
def get_info(self) -> dict:
"""返回索引信息(用于日志)"""
return {
'type': self.__class__.__name__,
'point_count': len(self.points),
'build_time': self.build_time
}

View File

@@ -0,0 +1,92 @@
"""Grid/Voxel 网格索引"""
import numpy as np
import time
from collections import defaultdict
from .base import SpatialIndex
class GridIndex(SpatialIndex):
"""Grid/Voxel 空间索引 - O(N)构建, O(1)查询"""
def __init__(self, points: np.ndarray, grid_size: float = 100.0):
"""
构建网格索引
Args:
points: 点云数组 (N, 3)
grid_size: 网格大小(毫米)
"""
super().__init__(points)
self.grid_size = grid_size
print(f" 构建Grid索引 (grid_size={grid_size}mm)...")
start = time.time()
# 计算点云边界
self.mins = points.min(axis=0)
self.maxs = points.max(axis=0)
# 将每个点映射到网格坐标
grid_coords = np.floor((points - self.mins) / grid_size).astype(np.int32)
# 构建索引:grid_key → point_indices
self.grid = defaultdict(list)
for idx, (gx, gy, gz) in enumerate(grid_coords):
grid_key = (gx, gy, gz)
self.grid[grid_key].append(idx)
# 转换为numpy数组以加速查询
for key in self.grid:
self.grid[key] = np.array(self.grid[key], dtype=np.int32)
self.build_time = time.time() - start
print(f" 网格数量: {len(self.grid):,}")
print(f" 平均点/格: {len(points) / len(self.grid):.1f}")
print(f" 构建耗时: {self.build_time:.2f} 秒")
def query_bbox(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
"""
Grid索引查询 - O(1 + K)复杂度
直接定位到相关网格,只检查这些网格中的点
"""
# 将bbox转换为网格坐标范围
grid_min = np.floor((bbox_min - self.mins) / self.grid_size).astype(np.int32)
grid_max = np.floor((bbox_max - self.mins) / self.grid_size).astype(np.int32)
# 收集相关网格中的所有点索引
candidate_indices = []
for gx in range(grid_min[0], grid_max[0] + 1):
for gy in range(grid_min[1], grid_max[1] + 1):
for gz in range(grid_min[2], grid_max[2] + 1):
grid_key = (gx, gy, gz)
if grid_key in self.grid:
candidate_indices.append(self.grid[grid_key])
if len(candidate_indices) == 0:
return np.array([], dtype=np.int32)
# 合并所有候选点
candidate_indices = np.concatenate(candidate_indices)
candidate_points = self.points[candidate_indices]
# 精确过滤(网格是粗略的,需要精确检查)
mask = (
(candidate_points[:, 0] >= bbox_min[0]) &
(candidate_points[:, 0] <= bbox_max[0]) &
(candidate_points[:, 1] >= bbox_min[1]) &
(candidate_points[:, 1] <= bbox_max[1]) &
(candidate_points[:, 2] >= bbox_min[2]) &
(candidate_points[:, 2] <= bbox_max[2])
)
return candidate_indices[mask]
def get_info(self) -> dict:
info = super().get_info()
info['method'] = 'Grid网格索引'
info['complexity'] = 'O(1 + K)'
info['grid_size'] = self.grid_size
info['grid_count'] = len(self.grid)
info['avg_points_per_grid'] = len(self.points) / len(self.grid)
return info

View File

@@ -0,0 +1,120 @@
"""Grid/Voxel 网格索引 - 优化版(numpy向量化)"""
import numpy as np
import time
from collections import defaultdict
from .base import SpatialIndex
class GridIndexOptimized(SpatialIndex):
"""Grid/Voxel 空间索引 - 优化版,使用numpy向量化避免Python循环"""
def __init__(self, points: np.ndarray, grid_size: float = 100.0):
"""
构建网格索引(优化版)
Args:
points: 点云数组 (N, 3)
grid_size: 网格大小(毫米)
"""
super().__init__(points)
self.grid_size = grid_size
print(f" 构建Grid索引(优化版 - numpy向量化)(grid_size={grid_size}mm)...")
start = time.time()
# 计算点云边界
self.mins = points.min(axis=0)
self.maxs = points.max(axis=0)
# 将每个点映射到网格坐标
grid_coords = np.floor((points - self.mins) / grid_size).astype(np.int32)
# ===== 关键优化:使用numpy向量化替代Python循环 =====
# 1. 将3D网格坐标编码为1D哈希值(避免tuple开销)
# 使用大素数来避免冲突
hash_multiplier = np.array([1, 1000003, 1000003**2], dtype=np.int64)
grid_hash = (grid_coords * hash_multiplier).sum(axis=1)
# 2. 使用argsort排序(numpy C实现,非常快)
sort_idx = np.argsort(grid_hash)
sorted_hash = grid_hash[sort_idx]
# 3. 找到每个唯一网格的边界
# 只需要遍历唯一网格数(约500万),而不是10亿点!
unique_mask = np.concatenate(([True], sorted_hash[1:] != sorted_hash[:-1]))
unique_indices = np.where(unique_mask)[0]
# 4. 构建索引字典
self.grid = {}
self.grid_hash_to_coords = {} # 保存hash到3D坐标的映射
n_unique = len(unique_indices)
for i in range(n_unique):
start_idx = unique_indices[i]
end_idx = unique_indices[i + 1] if i + 1 < n_unique else len(sorted_hash)
hash_val = sorted_hash[start_idx]
indices = sort_idx[start_idx:end_idx]
# 使用hash作为key(比tuple快)
self.grid[hash_val] = indices
# 保存hash到原始3D坐标的映射(用于查询)
self.grid_hash_to_coords[hash_val] = tuple(grid_coords[indices[0]])
self.build_time = time.time() - start
print(f" 网格数量: {len(self.grid):,}")
print(f" 平均点/格: {len(points) / len(self.grid):.1f}")
print(f" 构建耗时: {self.build_time:.2f} 秒")
def query_bbox(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
"""
Grid索引查询 - O(1 + K)复杂度
"""
# 将bbox转换为网格坐标范围
grid_min = np.floor((bbox_min - self.mins) / self.grid_size).astype(np.int32)
grid_max = np.floor((bbox_max - self.mins) / self.grid_size).astype(np.int32)
# 收集相关网格中的所有点索引
candidate_indices = []
# 遍历bbox覆盖的网格(通常很少,比如10×10×10=1000个)
hash_multiplier = np.array([1, 1000003, 1000003**2], dtype=np.int64)
for gx in range(grid_min[0], grid_max[0] + 1):
for gy in range(grid_min[1], grid_max[1] + 1):
for gz in range(grid_min[2], grid_max[2] + 1):
grid_coord = np.array([gx, gy, gz], dtype=np.int32)
grid_hash_val = (grid_coord * hash_multiplier).sum()
if grid_hash_val in self.grid:
candidate_indices.append(self.grid[grid_hash_val])
if len(candidate_indices) == 0:
return np.array([], dtype=np.int32)
# 合并所有候选点
candidate_indices = np.concatenate(candidate_indices)
candidate_points = self.points[candidate_indices]
# 精确过滤
mask = (
(candidate_points[:, 0] >= bbox_min[0]) &
(candidate_points[:, 0] <= bbox_max[0]) &
(candidate_points[:, 1] >= bbox_min[1]) &
(candidate_points[:, 1] <= bbox_max[1]) &
(candidate_points[:, 2] >= bbox_min[2]) &
(candidate_points[:, 2] <= bbox_max[2])
)
return candidate_indices[mask]
def get_info(self) -> dict:
info = super().get_info()
info['method'] = 'Grid网格索引(优化版)'
info['complexity'] = 'O(1 + K)'
info['grid_size'] = self.grid_size
info['grid_count'] = len(self.grid)
info['avg_points_per_grid'] = len(self.points) / len(self.grid)
info['optimization'] = 'numpy向量化'
return info

View File

@@ -0,0 +1,66 @@
"""KDTree 索引(使用scipy)"""
import numpy as np
import time
from scipy.spatial import cKDTree
from .base import SpatialIndex
class KDTreeIndex(SpatialIndex):
"""KDTree 空间索引 - O(N log N)构建, O(log N + K)查询"""
def __init__(self, points: np.ndarray):
"""
构建KD树索引
Args:
points: 点云数组 (N, 3)
"""
super().__init__(points)
print(f" 构建KDTree索引...")
start = time.time()
# 使用scipy的cKDTree(C实现,比纯Python快)
self.tree = cKDTree(points)
self.build_time = time.time() - start
print(f" 构建耗时: {self.build_time:.2f} 秒")
def query_bbox(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
"""
KDTree查询 - O(log N + K)复杂度
使用球形查询近似bbox,然后精确过滤
"""
# KDTree没有直接的bbox查询,使用球形查询近似
center = (bbox_min + bbox_max) / 2
# 使用外接球半径,确保包含整个bbox
radius = np.linalg.norm(bbox_max - bbox_min) / 2
# 球形查询
candidate_indices = self.tree.query_ball_point(center, radius)
if len(candidate_indices) == 0:
return np.array([], dtype=np.int32)
candidate_indices = np.array(candidate_indices, dtype=np.int32)
candidate_points = self.points[candidate_indices]
# 精确过滤(因为球形查询是近似的)
mask = (
(candidate_points[:, 0] >= bbox_min[0]) &
(candidate_points[:, 0] <= bbox_max[0]) &
(candidate_points[:, 1] >= bbox_min[1]) &
(candidate_points[:, 1] <= bbox_max[1]) &
(candidate_points[:, 2] >= bbox_min[2]) &
(candidate_points[:, 2] <= bbox_max[2])
)
return candidate_indices[mask]
def get_info(self) -> dict:
info = super().get_info()
info['method'] = 'KDTree索引(scipy)'
info['complexity'] = 'O(log N + K)'
info['note'] = '构建慢,适合小数据集'
return info

View File

@@ -0,0 +1,32 @@
"""线性扫描(无索引)- 基准方法"""
import numpy as np
import time
from .base import SpatialIndex
class LinearScanIndex(SpatialIndex):
"""线性扫描索引(实际上就是无索引)"""
def __init__(self, points: np.ndarray):
super().__init__(points)
# 无需构建索引
self.build_time = 0.0
def query_bbox(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
"""
线性扫描查询 - O(N)复杂度
遍历所有点,检查是否在bbox内
"""
mask = (
(self.points[:, 0] >= bbox_min[0]) & (self.points[:, 0] <= bbox_max[0]) &
(self.points[:, 1] >= bbox_min[1]) & (self.points[:, 1] <= bbox_max[1]) &
(self.points[:, 2] >= bbox_min[2]) & (self.points[:, 2] <= bbox_max[2])
)
return np.where(mask)[0]
def get_info(self) -> dict:
info = super().get_info()
info['method'] = '线性扫描(无索引)'
info['complexity'] = 'O(N)'
return info

View File

@@ -0,0 +1,166 @@
"""Octree 八叉树索引"""
import numpy as np
import time
from .base import SpatialIndex
class OctreeNode:
"""八叉树节点"""
def __init__(self, center: np.ndarray, size: float):
self.center = center
self.size = size
self.children = [None] * 8 # 8个子节点
self.point_indices = [] # 叶节点存储点索引
self.is_leaf = True
class OctreeIndex(SpatialIndex):
"""Octree 空间索引 - O(N)构建, O(log N + K)查询"""
def __init__(self, points: np.ndarray, max_depth: int = 8, max_points_per_node: int = 1000):
"""
构建八叉树索引
Args:
points: 点云数组 (N, 3)
max_depth: 最大树深度
max_points_per_node: 每个叶节点最多存储的点数
"""
super().__init__(points)
self.max_depth = max_depth
self.max_points_per_node = max_points_per_node
print(f" 构建Octree索引 (max_depth={max_depth}, max_points={max_points_per_node})...")
start = time.time()
# 计算点云边界
mins = points.min(axis=0)
maxs = points.max(axis=0)
center = (mins + maxs) / 2
size = np.max(maxs - mins) / 2
# 构建树
all_indices = np.arange(len(points), dtype=np.int32)
self.root = self._build_tree(center, size, all_indices, depth=0)
self.build_time = time.time() - start
self.node_count = self._count_nodes(self.root)
print(f" 节点数量: {self.node_count:,}")
print(f" 构建耗时: {self.build_time:.2f} 秒")
def _build_tree(self, center: np.ndarray, size: float, indices: np.ndarray, depth: int) -> OctreeNode:
"""递归构建八叉树"""
node = OctreeNode(center, size)
# 终止条件:达到最大深度或点数足够少
if depth >= self.max_depth or len(indices) <= self.max_points_per_node:
node.point_indices = indices
node.is_leaf = True
return node
# 分割成8个子空间
node.is_leaf = False
half_size = size / 2
points_subset = self.points[indices]
for i in range(8):
# 计算子节点中心
child_center = center.copy()
child_center[0] += half_size if (i & 1) else -half_size
child_center[1] += half_size if (i & 2) else -half_size
child_center[2] += half_size if (i & 4) else -half_size
# 找到属于这个子空间的点
mask = (
((points_subset[:, 0] >= child_center[0] - half_size) &
(points_subset[:, 0] < child_center[0] + half_size)) &
((points_subset[:, 1] >= child_center[1] - half_size) &
(points_subset[:, 1] < child_center[1] + half_size)) &
((points_subset[:, 2] >= child_center[2] - half_size) &
(points_subset[:, 2] < child_center[2] + half_size))
)
child_indices = indices[mask]
if len(child_indices) > 0:
node.children[i] = self._build_tree(child_center, half_size, child_indices, depth + 1)
return node
def query_bbox(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
"""
Octree查询 - O(log N + K)复杂度
递归查询与bbox相交的节点
"""
result = []
self._query_recursive(self.root, bbox_min, bbox_max, result)
if len(result) == 0:
return np.array([], dtype=np.int32)
# 合并所有候选点
candidate_indices = np.concatenate(result)
candidate_points = self.points[candidate_indices]
# 精确过滤
mask = (
(candidate_points[:, 0] >= bbox_min[0]) &
(candidate_points[:, 0] <= bbox_max[0]) &
(candidate_points[:, 1] >= bbox_min[1]) &
(candidate_points[:, 1] <= bbox_max[1]) &
(candidate_points[:, 2] >= bbox_min[2]) &
(candidate_points[:, 2] <= bbox_max[2])
)
return candidate_indices[mask]
def _query_recursive(self, node: OctreeNode, bbox_min: np.ndarray, bbox_max: np.ndarray, result: list):
"""递归查询节点"""
if node is None:
return
# 检查节点边界框是否与查询bbox相交
node_min = node.center - node.size
node_max = node.center + node.size
if not self._bbox_intersects(bbox_min, bbox_max, node_min, node_max):
return
if node.is_leaf:
# 叶节点:添加所有点索引
if len(node.point_indices) > 0:
result.append(node.point_indices)
else:
# 内部节点:递归查询子节点
for child in node.children:
if child is not None:
self._query_recursive(child, bbox_min, bbox_max, result)
def _bbox_intersects(self, bbox1_min: np.ndarray, bbox1_max: np.ndarray,
bbox2_min: np.ndarray, bbox2_max: np.ndarray) -> bool:
"""检查两个bbox是否相交"""
return not (
bbox1_max[0] < bbox2_min[0] or bbox1_min[0] > bbox2_max[0] or
bbox1_max[1] < bbox2_min[1] or bbox1_min[1] > bbox2_max[1] or
bbox1_max[2] < bbox2_min[2] or bbox1_min[2] > bbox2_max[2]
)
def _count_nodes(self, node: OctreeNode) -> int:
"""统计节点数量"""
if node is None:
return 0
count = 1
if not node.is_leaf:
for child in node.children:
count += self._count_nodes(child)
return count
def get_info(self) -> dict:
info = super().get_info()
info['method'] = 'Octree八叉树索引'
info['complexity'] = 'O(log N + K)'
info['max_depth'] = self.max_depth
info['node_count'] = self.node_count
return info