feat: update algorithm integrations and point cloud overlays
This commit is contained in:
@@ -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',
|
||||
]
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user