FarthestPointSampling(FPS)算法学习&实

最近工作需要用到最远点采样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