Function gather_points

pointersect/pr/naive.py:159–189

gather the points belonging to each grid cell. Args: grid_idxs: grid index of each point total_cells: total grid cells valid_mask: whether to record the point Returns: list of list of list, (b, total_cells[bidx]) -> n

(
        grid_idxs: torch.Tensor,  # (b, n), long
        total_cells: torch.Tensor,  # (b,) long
        valid_mask: torch.Tensor = None,  # (b, n), bool
)

def gather_points(
        grid_idxs: torch.Tensor,  # (b, n), long
        total_cells: torch.Tensor,  # (b,) long
        valid_mask: torch.Tensor = None,  # (b, n), bool
) -> T.List[T.List[T.List[int]]]:
    """
    gather the points belonging to each grid cell.

    Args:
        grid_idxs:
            grid index of each point
        total_cells:
            total grid cells
        valid_mask:
            whether to record the point

    Returns:
        list of list of list, (b, total_cells[bidx]) -> n_idx of points in the cell
    """
    batch_size, n_points = grid_idxs.shape
    grid_idxs = grid_idxs.long()
    total_cells = total_cells.long()

    all_cell2pidx = []
    for b in range(batch_size):
        cell2pidx: T.List[T.List[int]] = [[] for _ in range(total_cells[b])]
        for i in range(n_points):
            if valid_mask is None or valid_mask[b, i]:
                cell2pidx[grid_idxs[b, i]].append(i)
        all_cell2pidx.append(cell2pidx)
    return all_cell2pidx


def grid_ray_intersection(