Added functionality for point cloud with pytorch tensors and PCA analysis, all working on CUDA

[8ByteSword]

For my previous job, I got a project for detection of vegetation encroachment for powerlines using LiDAR point clouds, and I end up needing to implement my own functions over the open3d.t.geometry.PointCloud class for better working with pytorch tensors as it was absolutelly necessary to perform compute on GPU. I was wondering who to contact or ask to see if it is a valuable addition to open3d or if I should make my own repo for this (while aknowledging open3d as it was just overloaded). I made functions to work directly with pytorch tensors for many functions like crop, raster, select by index, clone, ... Also implemented a VoxelGrid class for neighborhood dependant operations, which maintains all the data in cuda for same reasons. On top of that, I developed my own functions for computing neighborhood based eigenvalues calculation, PCA analysis, catenary fitting, also using faiss search for both voxel neighborhood and radius search.
Who should I talk with?

A particular use case that uses my VoxelGrid class, generated using my custom PointCloud class that extends o3d.t.geometry.Pointcloud, as example:

def compute_eigenvalues(k: int, radius: float, voxel_grid: VoxelGrid, batch_size: int = 10000) -> Tensor:
    """
    Compute eigenvalues for the point cloud using batched processing.

    Args:
        k (int): Number of nearest neighbors to consider.
        radius (float): Radius for filtering.
        voxel_grid (VoxelGrid): Voxel grid of the point cloud.
        batch_size (int): Size of batches for processing.

    Returns:
        Tensor: Tensor of computed eigenvalues.
    """
    d = 3
    index = faiss.IndexFlatL2(d)
    res = faiss.StandardGpuResources()
    gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
    eigenvalues_tensor = torch.zeros(
        (voxel_grid.pcd.positions.shape[0], 3),
        device='cuda',
        dtype=torch.float64
    )

    points_shifted = voxel_grid.pcd.positions - voxel_grid.min_coords

    batches = []
    for central_linear_index, neighbor_voxel_indices in voxel_grid.precomputed_neighbor_voxels.items():
        valid_mask = torch.isin(neighbor_voxel_indices, voxel_grid.unique)
        valid_voxel_indices = neighbor_voxel_indices[valid_mask]
        batches.append((central_linear_index, valid_voxel_indices))
        if len(batches) == batch_size:
            process_batch(batches, points_shifted, gpu_index, k, radius, voxel_grid, eigenvalues_tensor)
            batches = []

    if batches:
        process_batch(batches, points_shifted, gpu_index, k, radius, voxel_grid, eigenvalues_tensor)

    return eigenvalues_tensor

def process_batch(batches: List[Tuple[int, Tensor]], points_shifted: Tensor, gpu_index: faiss.GpuIndex, 
                  k: int, radius: float, voxel_grid: VoxelGrid, eigenvalues_tensor: Tensor) -> None:
    """
    Process a batch of voxels to compute eigenvalues.

    Args:
        batches (List[Tuple[int, Tensor]]): List of batches to process.
        points_shifted (Tensor): Shifted point cloud positions.
        gpu_index (faiss.GpuIndex): GPU index for nearest neighbor search.
        k (int): Number of nearest neighbors to consider.
        radius (float): Radius for filtering.
        voxel_grid (VoxelGrid): Voxel grid of the point cloud.
        eigenvalues_tensor (Tensor): Tensor to store computed eigenvalues.
    """
    try:
        logger.info("Processing batch...")
        for central_linear_index, valid_voxel_indices in batches:
            voxel_points = points_shifted[voxel_grid.voxel_to_indices[central_linear_index]]
            neighbor_points = torch.cat(
                [points_shifted[voxel_grid.voxel_to_indices[voxel_index.item()]] for voxel_index in valid_voxel_indices]
            )
            neighbor_points = neighbor_points.to(torch.float32)
            voxel_points = voxel_points.to(torch.float32)

            gpu_index.add(neighbor_points) # type: ignore
            distances, indices = gpu_index.search(voxel_points, min(k, neighbor_points.shape[0])) # type: ignore

            centered_neighbors, num_points = radius_filter_and_recenter(neighbor_points, distances, indices, radius)

            cov_matrices = torch.matmul(centered_neighbors.transpose(1, 2), centered_neighbors) / num_points.repeat_interleave(3).view(-1, 3, 3)

            eigenvalues = torch.linalg.eigvalsh(cov_matrices).to(torch.float64) 
            original_indices = voxel_grid.voxel_to_indices[central_linear_index]
            eigenvalues_tensor[original_indices] = eigenvalues

            gpu_index.reset()
    except RuntimeError as e:
        logger.info(f"CUDA error encountered: {e}")
        # Implement fallback or error handling strategy here