最近工作需要用到最远点采样fps算法, 找了一些blog学习学习, 写了个原型:
def distance(p1, p2):
return np.sqrt(np.sum((p1 - p2)**2))
def FPS(sample, num, center_idx: int = 0):
'''sample:采样点云数据,
num:需要采样的数据点个数'''
n = sample.shape[0]
select_p = [center_idx] # 储存采集点索引
L = np.full(n, 1e20, dtype=np.float64)
for i in range(num - 1):
for p in range(n):
d = distance(sample[select_p[-1]], sample[p])
if d < L[p]:
L[p] = d
select_p.append(np.argmax(L).item())
return select_p
测试一下, 跑的没问题, 但是速度感人
所有点的欧氏距离计算做个向量化优化:
def fps(points, n_samples, first_idx: int = 0):
"""
points: (N, D) array of points
n_samples: number of points to sample
first_idx: index of first point to sample
"""
N = points.shape[0]
selected_indices = [first_idx]
distances = np.full(N, np.inf)
for _ in range(n_samples - 1):
# 计算所有点到最新选择点的距离
last_selected = points[selected_indices[-1]]
dist = np.sum((points - last_selected) ** 2, axis=1)
# 更新最小距离
distances = np.minimum(distances, dist)
# 选择距离最大的点
new_idx = np.argmax(distances).item()
selected_indices.append(new_idx)
distances[new_idx] = 0
return selected_indices
这一下子快了100多倍, 能用了
找了个开源基于rust实现核心的python接口包: https://github.com/leonardodalinky/fpsample
对比了一下结果, 是没问题的, 但是他比我这个快个4,5倍
最后借助ai优化一下:
def fps_optimized(points, n_samples, first_idx: int = 0):
"""
Optimized FPS implementation
points: (N, D) array of points
n_samples: number of points to sample
first_idx: index of first point to sample
"""
N = points.shape[0]
selected_indices = np.zeros(n_samples, dtype=np.int64)
distances = np.full(N, np.inf)
selected_indices[0] = first_idx
# 预计算点的平方和,避免重复计算
point_squares = np.sum(points ** 2, axis=1)
for i in range(1, n_samples):
last_selected = points[selected_indices[i - 1]]
# 使用广播来计算距离,避免创建临时数组
# ||a-b||^2 = ||a||^2 + ||b||^2 - 2<a,b>
dist = point_squares + np.sum(last_selected ** 2) - 2 * np.dot(points, last_selected)
# 更新最小距离
np.minimum(distances, dist, out=distances)
# 选择距离最大的点
new_idx = np.argmax(distances)
selected_indices[i] = new_idx
distances[new_idx] = 0
return selected_indices.tolist()
又快了6,7倍, 中小规模点云用这个看起来就没啥问题了
后面如果有超大规模的需求或许可以关注fpsample作者的另一个repo: https://github.com/leonardodalinky/pytorch_fpsample