Pointnet++ code explanation (IV): index_ Points function

index_points The function returns the indexed point cloud data according to the input point cloud data and index . for example points by B×2048×3 Point cloud of ,idx by [1,333,1000,2000](S Dimensions ), Then return to B Number of samples in each sample 1,333,1000,2000 It's made up of dots B×4×3 A cloud of points . Of course, if idx For one [B,S] Dimensional , Then it will follow idx The dimension structure in is extracted into [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], namely 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 Dimensions [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

Yes points[batch_indices, idx, :] The understanding of the ：

batch_indices yes tensor, Dimension for [B,S], It represents the sample index , That is, which sample to take , Its expression is as follows ：

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

idx It's also tensor, The dimension is also [B,S], It means that each line corresponds to 1 Samples , The data in this row is in points Index of data in .

points[batch_indices, idx, :] That's from points Take out each batch The data point corresponding to the index .