K近邻算法（KNN）和最远点采样（FPS）实现--python+pytorch

最新推荐文章于 2023-09-13 07:49:03 发布

Tony

最新推荐文章于 2023-09-13 07:49:03 发布

阅读量4k

点赞数

分类专栏：深度学习文章标签： pytorch python 算法

本文链接：https://blog.csdn.net/qq_44643484/article/details/123638855

版权

深度学习专栏收录该内容

16 篇文章 2 订阅

订阅专栏

本文介绍了如何使用Python和PyTorch实现K近邻（KNN）算法，以及最远点采样（FPS）方法。首先，通过生成点集，展示了KNN算法的详细步骤，计算每个点到某点的距离并找到最近的K个点。接着，给出了一个完整的KNN示例，包括计算距离和寻找近邻点下标。此外，还提供了远点采样的实现，通过不断选取最远的点进行采样，生成指定数量的采样点集。

摘要由CSDN通过智能技术生成

K近邻算法（KNN）

主要思路：计算每个点和某点的距离，取距离最短的K个点的下标即可。
下面是个完整示例，代码复制即可运行

import torch
import time

#生成点集
def coordinate_gen(n):
    """
    生成n个三位点
    return tensor
    dim:n*3
    """
    xyz = torch.rand(size=(n,3))
    return xyz

#计算运行时间
def time_cost(f):

    def run_time(*args,**kwargs):
        start = time.time()
        res = f(*args,**kwargs)
        run_times = time.time()-start
        print("程序执行时间：%.6f s"%(run_times))
        return res

    return run_time

@time_cost
def knn(xyz,xyzs,k=3):
    """
    xyz:key point
    xyzs:all points
    找某点的k近邻个点
    return 近邻点的下标列表
    """
    idx = [0]*k
    distance = torch.sum((xyzs[:,:3]-xyz[:,:3])**2,dim=-1)
    for i in range(k):
        idx[i] = torch.argmin(distance,dim=0)
        distance[int(torch.argmin(distance,dim=0))] = float('inf')
    idx = [int(i) for i in idx]
    return idx

if __name__ == "__main__":
    print('-' * 20, '测试开始', '-' * 20)
    N,k = map(int,input("输入生成点数 和 k的值:").split())
    xyzs = coordinate_gen(N)
    xyz = torch.rand(size=(1,3))
    print("生成的点如下:\n",xyzs,"\n随机生成key point:",xyz)
    print(knn(xyz,xyzs=xyzs,k=k))
    print('-'*20,'测试结束','-'*20)

最远点采样（FPS）

def farthest_point_sample(data,npoints):
    """
    Args:
        data:输入的tensor张量，排列顺序 N,D
        Npoints: 需要的采样点

    Returns:data->采样点集组成的tensor，每行是一个采样点
    """
    N,D = data.shape #N是点数，D是维度
    xyz = data[:,:3] #只需要坐标
    centroids = torch.zeros(size=(npoints,)) #最终的采样点index
    dictance = torch.ones(size=(N,))*1e10 #距离列表,一开始设置的足够大,保证第一轮肯定能更新dictance
    farthest = torch.randint(low=0,high=N,size=(1,)) #随机选一个采样点的index
    for i in range(npoints):
        centroids[i] = farthest
        centroid = xyz[farthest,:]
        dict = ((xyz-centroid)**2).sum(dim=-1)
        mask = dict < dictance
        dictance[mask] = dict[mask]
        farthest = torch.argmax(dictance,dim=-1)
    print(centroids.type(torch.long))
    data= data[centroids.type(torch.long)]
    return data