def kdtree_radius_search(root: Node, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 作业3
# 提示：通过递归的方式实现搜索
# 屏蔽开始
if query[root.axis] <= root.value:
kdtree_radius_search(root.left, db, result_set, query)
if math.fabs(query[root.axis] - root.value) < result_set.worstDist():
kdtree_radius_search(root.right, db, result_set, query)
else:
kdtree_radius_search(root.right, db, result_set, query)
if math.fabs(query[root.axis] - root.value) < result_set.worstDist():
kdtree_radius_search(root.left, db, result_set, query)
# 屏蔽结束
return False
def kdtree_radius_search(root: Node, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 作业3
# 提示：通过递归的方式实现搜索
# 屏蔽开始
if root.value >= query[root.axis]:
if kdtree_knn_search(root.left, db, result_set, query):
return True
elif math.fabs(root.value - query[root.axis]) < result_set.worst_dist:
return kdtree_knn_search(root.right, db, result_set, query)
return False # 3、左右子树都不满足，就需要返回上一层树，所以为False
else:
if kdtree_knn_search(root.right, db, result_set, query):
return True
elif math.fabs(root.value - query[root.axis]) < result_set.worst_dist:
return kdtree_knn_search(root.left, db, result_set, query)
return False # 3、左右子树都不满足，就需要返回上一层树，所以为False
# 屏蔽结束
return False
def octree_radius_search_fast(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if contains(query, result_set.worstDist(), root):
# compare the contents of the octant
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# don't need to check any child
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# no need to go to most relevant child first, because anyway we will go through all children
for c, child in enumerate(root.children):
if child is None:
continue
if False == overlaps(query, result_set.worstDist(), child):
continue
if octree_radius_search_fast(child, db, result_set, query):
return True
return inside(query, result_set.worstDist(), root)
def octree_radius_search_fast(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
# 1. if query ball contains the octant, no need to check child, just compare all point in it!
# 只是不需要向下递归，这里的child是只当前节点的下一层，回溯还是有可能的！所以这里return的false是有原因的嗷！
if contains(query, result_set.worstDist(), root):
# compare all points:
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# no need to check child
return False
# consider leaf point:
if root.is_leaf and len(root.point_indices) > 0:
# compare all points:
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return inside(query, result_set.worstDist(), root)
# 2. check all children
for c, child in enumerate(root.children):
if child is None:
continue
if False == overlaps(query, result_set.worstDist(), child):
continue
if octree_radius_search_fast(child, db, result_set,query):
return True
return inside(query, result_set.worstDist(), root)
def octree_radius_search(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 作业6
# 屏蔽开始
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# go to the relevant child first
morton_code = 0
if query[0] > root.center[0]:
morton_code = morton_code | 1
if query[1] > root.center[1]:
morton_code = morton_code | 2
if query[2] > root.center[2]:
morton_code = morton_code | 4
if octree_radius_search(root.children[morton_code], db, result_set, query):
return True
# check other children
for c, child in enumerate(root.children):
if c == morton_code or child is None:
continue
if False == overlaps(query, result_set.worstDist(), child):
continue
if octree_radius_search(child, db, result_set, query):
return True
# 屏蔽结束
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def octree_radius_search_plus(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 提前结束：核心-八叉树对3个维度有限制；
# 结束条件：其实是两个方向：向下递归，向上回溯
# 1-当一个节点返回True时，表示找到knn了，维度限制导致不会有更优的了，不再向下，立即结束；
# 2-最坏距离球如果 inside 当前节点，不需要再向上回溯检查其他节点，立即结束；
# 跳过条件：在检查
# 1. search the first relevant child:
# 找到最近的孩子，根据查询点的莫顿码
morton_code = 0
if query[0] > root.center[0]:
morton_code = morton_code | 1
if query[1] > root.center[1]:
morton_code = morton_code | 2
if query[2] > root.center[2]:
morton_code = morton_code | 4
if octree_knn_search(root.children[morton_code],
db,
result_set,
query):
return True
# 2. check other children
inv_worstDist = 1 / result_set.worstDist()
for c, child in enumerate(root.children):
if c == morton_code or child is None:
continue
if False == overlaps_plus(query, inv_worstDist, child):
continue
if octree_knn_search(child, db, result_set, query):
return True
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def octree_radius_search_fast(root: Octant, db: np.ndarray,
result_set: RadiusNNResultSet,
query: np.ndarray):
if root is None:
return False
# 作业5
# 提示：尽量利用上面的inside、overlaps、contains等函数
# 屏蔽开始 #如果球包含了这个root Octant则不需要去它的子节点找了
if (root.is_leaf and len(root.point_indices) > 0) or contains(
query, result_set.worstDist(), root):
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
#如果球包含了这个root Octant则不需要去它的子节点找了
children_idx = 0
if query[0] > root.center[0]: # x轴
children_idx = children_idx | 1
if query[1] > root.center[1]: # y轴
children_idx = children_idx | 2
if query[2] > root.center[2]: # z轴
children_idx = children_idx | 4
# 如果在这个子octant中发现，查询点在该子octant中，同时由worst_dist构成的球也被octant包含，所以直接返回
if octree_knn_search(root.children[children_idx], db, result_set, query):
return True
# 如果不满足上边的情况则需要遍历其他子octant
for c, child in enumerate(root.children):
if c == children_idx or root.children[c]:
continue
if overlaps(query, result_set.worstDist(), child) == False:
continue
if octree_knn_search(root.children[c], db, result_set, query):
return True
# 屏蔽结束
return inside(query, result_set.worstDist(), root)
def octree_radius_search(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 作业6
# 屏蔽开始
# 屏蔽结束
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def kdtree_radius_search(root: Node, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 作业3
# 提示：通过递归的方式实现搜索
# 屏蔽开始
# 屏蔽结束
return False
def radius_search(root: Node, result_set: RadiusNNResultSet, key):
if root is None:
return False
# compare the root itself
result_set.add_point(math.fabs(root.key - key), root.value)
if root.key >= key:
# iterate left branch first
if radius_search(root.left, result_set, key):
return True
elif math.fabs(root.key-key) < result_set.worstDist():
return radius_search(root.right, result_set, key)
return False
else:
# iterate right branch first
if radius_search(root.right, result_set, key):
return True
elif math.fabs(root.key-key) < result_set.worstDist():
return radius_search(root.left, result_set, key)
return False
def octree_radius_search_fast(root: Octant, db: np.ndarray,
result_set: RadiusNNResultSet,
query: np.ndarray):
if root is None:
return False
# 作业5
# 提示：尽量利用上面的inside、overlaps、contains等函数
# 屏蔽开始
if contains(query, result_set.worstDist(), root):
# the octant is contained by the query ball with resule_set.worstDist()
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# check other children
for c, child in enumerate(root.children):
if child is None:
continue
# if an octant is not overlapping with query ball, skip
if overlaps(query, result_set.worstDist(), child) == False:
continue
if octree_radius_search(child, db, result_set, query):
return True
# 屏蔽结束
return inside(query, result_set.worstDist(), root)
def kdtree_radius_search(root: Node, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return
if root.is_leaf():
leaf_points = db[root.point_indices]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return
if query[root.axis] <= root.value:
kdtree_radius_search(root.left, db, result_set, query)
if math.fabs(query[root.axis] -
root.value) < result_set.get_worst_dist():
kdtree_radius_search(root.right, db, result_set, query)
else:
kdtree_radius_search(root.right, db, result_set, query)
if math.fabs(query[root.axis] -
root.value) < result_set.get_worst_dist():
kdtree_radius_search(root.left, db, result_set, query)
def main():
# configuration
leaf_size = 32
min_extent = 0.0001
k = 8
radius = 1
#root_dir = '/Users/renqian/cloud_lesson/kitti' # 数据集路径
#cat = os.listdir(root_dir)
#iteration_num = len(cat)
# load date
filename = '/home/ljn/SLAM/dateset/000000.bin'
db_np = read_velodyne_bin(filename)
iteration = 1
print("octree -----------------------------------")
construction_time_sum = 0
knn_time_sum = 0
radius_time_sum = 0
brute_time_sum = 0
result_set_knn = KNNResultSet(capacity=k)
query = db_np[95, :]
begin_t = time.time()
root = octree.octree_construction(db_np, leaf_size, min_extent)
construction_time_sum += time.time() - begin_t
begin_t = time.time()
octree.octree_knn_search(root, db_np, result_set_knn, query)
knn_time_sum += time.time() - begin_t
print('knn search result\n', result_set_knn)
begin_t = time.time()
result_set_rnn = RadiusNNResultSet(radius=radius)
octree.octree_radius_search_fast(root, db_np, result_set_rnn, query)
radius_time_sum += time.time() - begin_t
print('rnn search result\n', result_set_rnn)
begin_t = time.time()
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
nn_dist_idx_pre = np.linspace(0, nn_dist.shape[0], nn_dist.shape[0])
nn_dist_idx = nn_dist_idx_pre[nn_idx]
brute_time_sum += time.time() - begin_t
#brute knn search
result_set_knn_brute = KNNResultSet(capacity=k)
for index in range(k):
result_set_knn_brute.add_point(nn_dist[index], nn_dist_idx[index])
# brute radiusNN search
result_set_rnn_brute = RadiusNNResultSet(radius=radius)
for index in range(nn_dist.shape[0]):
if nn_dist[index] < radius:
result_set_rnn_brute.add_point(nn_dist[index], nn_dist_idx[index])
continue
else:
break
print("Octree: build %.3f, knn %.3f, radius %.3f, brute %.3f" %
(construction_time_sum * 1000, knn_time_sum * 1000,
radius_time_sum * 1000, brute_time_sum * 1000))
print("kdtree -----------------------------------")
construction_time_sum = 0
knn_time_sum = 0
radius_time_sum = 0
brute_time_sum = 0
for i in range(iteration):
query = db_np[95, :]
begin_t = time.time()
root = kdtree.kdtree_construction(db_np, leaf_size)
construction_time_sum += time.time() - begin_t
begin_t = time.time()
result_set_knn = KNNResultSet(capacity=k)
kdtree.kdtree_knn_search(root, db_np, result_set_knn, query)
knn_time_sum += time.time() - begin_t
begin_t = time.time()
result_set_rnn = RadiusNNResultSet(radius=radius)
kdtree.kdtree_radius_search(root, db_np, result_set_rnn, query)
radius_time_sum += time.time() - begin_t
begin_t = time.time()
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
brute_time_sum += time.time() - begin_t
nn_dist_idx_pre = np.linspace(0, nn_dist.shape[0] - 1,
nn_dist.shape[0])
nn_dist_idx = nn_dist_idx_pre[nn_idx]
# brute knn search
result_set_knn_brute = KNNResultSet(capacity=k)
for index in range(k):
result_set_knn_brute.add_point(nn_dist[index], nn_dist_idx[index])
# brute radiusNN search
result_set_rnn_brute = RadiusNNResultSet(radius=radius)
for index in range(nn_dist.shape[0]):
if nn_dist[index] < radius:
result_set_rnn_brute.add_point(nn_dist[index],
nn_dist_idx[index])
continue
else:
break
print("Kdtree: build %.3f, knn %.3f, radius %.3f, brute %.3f" %
(construction_time_sum * 1000, knn_time_sum * 1000,
radius_time_sum * 1000, brute_time_sum * 1000))
print("scipy kdtree -----------------------------------")
construction_time_sum = 0
knn_time_sum = 0
radius_time_sum = 0
query = db_np[95, :]
begin_t = time.time()
tree = spatial.KDTree(db_np, leaf_size)
construction_time_sum += time.time() - begin_t
#no knn
begin_t = time.time()
knn_time_sum += time.time() - begin_t
begin_t = time.time()
result_set_rnn = tree.query_ball_point(query, radius)
radius_time_sum += time.time() - begin_t
print('rnn search result\n', result_set_rnn)
print("Octree: build %.3f, knn %.3f, radius %.3f" %
(construction_time_sum * 1000, knn_time_sum * 1000,
radius_time_sum * 1000))
def octree_radius_search(root: Octant,
db: np.ndarray,
result_set: RadiusNNResultSet,
query: np.ndarray,
search='normal'):
if root is None:
return False
# 作业5
# 另外一种极端情况：查询半径已经包括了整个区域
if search == 'fast':
if contains(query, result_set.worstDist(), root):
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points,
axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 如果是叶子，则将该区域内返回到result_set
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 作业6
# 屏蔽开始
child_idx = -1
if search == 'normal':
child_idx = 0
# 判断所查询点在八叉树的位置
if query[0] > root.center[0]:
child_idx = child_idx | 1
if query[1] > root.center[1]:
child_idx = child_idx | 2
if query[2] > root.center[2]:
child_idx = child_idx | 4
# 递归判断是否在该区域就可以找到足够多的点
if octree_radius_search(root.children[child_idx], db, result_set,
query, search):
return True
# 没有在查询点区域找到，对其他区域进行搜索
for c, child in enumerate(root.children):
if c == child_idx or child is None: # 搜索区域没有点或所搜索到查询点区域，skip
continue
if False == overlaps(query, result_set.worstDist(),
child): # 搜索区域与查询点和最坏距离构成的球面没有交点，skip
continue
if octree_radius_search(child, db, result_set, query,
search): # 其他情况可以进入搜索区域搜索（递归）
return True
# 屏蔽结束
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
