Pointnet++代码详解（四）：index_points函数

index_points函数是按照输入的点云数据和索引返回由索引的点云数据。例如points为B×2048×3的点云，idx为[1,333,1000,2000](S的维度)，则返回B个样本中每个样本的第1,333,1000,2000个点组成的B×4×3的点云集。当然如果idx为一个[B,S]维度的，则它会按照idx中的维度结构将其提取成[B,S,C]。

def index_points(points, idx):
    """
    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    batchsize = points.shape[0]
    #view_shape=[B,S]
    view_shape = list(idx.shape)
    #[1] * (len(view_shape) - 1) -> [1],即view_shape=[B,1]
    view_shape[1:] = [1] * (len(view_shape) - 1)
    #repeat_shape=[B,S]
    repeat_shape = list(idx.shape)
    #repeat_shape=[1,S]
    repeat_shape[0] = 1
    #.view(view_shape)=.view(B,1)
    #.repeat(repeat_shape)=.view(1,S)
    #batch_indices的维度[B,S]
    batch_indices = torch.arange(batchsize, dtype=torch.long).to(device)\
        .view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points

对points[batch_indices, idx, :]的理解：

batch_indices是tensor,维度为[B,S]，表示的是样本索引，即取哪个样本，其表示形式如下：

tensor([[0, 0, 0, 0,...],
        [1, 1, 1, 1,...],...,
        [B-1, B-1, B-1, B-1,...]])

idx也是tensor,维度也为[B,S]，其表示的是每一行对应1个样本，该行中的数据就是在points当中数据的索引。

points[batch_indices, idx, :]说的就是从points当中取出每个batch对应索引的数据点。