PointNet++代码的实现

代码主要由两部分组成，pointnet_util.py封装着一些重要的函数组件，pointnet2.py用来搭建模型。

1. 功能函数文件

1、square_distance函数
该函数主要用来在ball query过程中确定每一个点距离采样点的距离。函数输入是两组点，N为第一组点的个数，M为第二组点的个数，C为输入点的通道数（如果是xyz时C=3），返回的是两组点之间两两的欧几里德距离，即N × M N\times MN×M的矩阵。由于在训练中数据通常是以Mini-Batch的形式输入的，所以有一个Batch数量的维度为B。

def square_distance(src, dst):
    """
    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2     
	     = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) # 2*(xn * xm + yn * ym + zn * zm)
    dist += torch.sum(src ** 2, -1).view(B, N, 1) # xn*xn + yn*yn + zn*zn
    dist += torch.sum(dst ** 2, -1).view(B, 1, M) # xm*xm + ym*ym + zm*zm
    return dist

2、 farthest_point_sample函数
最远点采样是Set Abstraction模块中较为核心的步骤，其目的是从一个输入点云中按照所需要的点的个数npoint采样出足够多的点，并且点与点之间的距离要足够远。最后的返回结果是npoint个采样点在原始点云中的索引。
具体步骤如下：

• 先随机初始化一个centroids矩阵，后面用于存储npoint个采样点的索引位置，大小为B × n p o i n t B\times
  npointB×npoint，其中B为BatchSize的个数，即B个样本;
• 利用distance矩阵记录某个样本中所有点到某一个点的距离，初始化为B×N矩阵，初值给个比较大的值，后面会迭代更新;
• 利用farthest表示当前最远的点，也是随机初始化，范围为0~N，初始化B个，对应到每个样本都随机有一个初始最远点；
• batch_indices初始化为0~(B-1)的数组；
• 直到采样点达到npoint，否则进行如下迭代：
  • 设当前的采样点centroids为当前的最远点farthest；
  • 取出这个中心点centroid的坐标； 求出所有点到这个farthest点的欧式距离，存在dist矩阵中；
  • 建立一个mask，如果dist中的元素小于distance矩阵中保存的距离值，则更新distance中的对应值，随着迭代的继续distance矩阵中的值会慢慢变小，其相当于记录着某个样本中每个点距离所有已出现的采样点的最小距离；
  • 最后从distance矩阵取出最远的点为farthest，继续下一轮迭代

def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, C]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
    	# 更新第i个最远点
        centroids[:, i] = farthest
        # 取出这个最远点的xyz坐标
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        # 计算点集中的所有点到这个最远点的欧式距离
        dist = torch.sum((xyz - centroid) ** 2, -1)
        # 更新distances，记录样本中每个点距离所有已出现的采样点的最小距离
        mask = dist < distance
        distance[mask] = dist[mask]
        # 从更新后的distances矩阵中找出距离最远的点，作为最远点用于下一轮迭代
        farthest = torch.max(distance, -1)[1]
    return centroids

标签：distance,npoint,torch,dist,farthest,++,代码,PointNet,xyz
来源： https://blog.csdn.net/weixin_45947476/article/details/121096056