isl-org / Open3D-ML

An extension of Open3D to address 3D Machine Learning tasks
Other
1.83k stars 313 forks source link

Query points in neighbor search FixedRadiusSearch #491

Open PabloVD opened 2 years ago

PabloVD commented 2 years ago

Checklist

My Question

Hi, I want to generate a graph from a cloud of points, getting all edges connecting points closer than a given distance. For that, the best option seems to be torch.layers.FixedRadiusSearch, which gives neighbors_index as the indexes of neighbors. However, it does not output the query points corresponding to these neighbors. The only mention about that appears in the neighbors_index documentation:

The compact list of indices of the neighbors. The corresponding query point can be inferred from the ‘neighbor_count_row_splits’ vector.

The mentioned neighbor_count_row_splits vector seems to be undefined and does not appear in the documentation.

Is there a way to efficiently obtain the query points from the FixedRadiusSearch output? Many thanks in advance!

biophase commented 2 years ago

Not sure I can help you much as I have only done this with the tensorflow equivalent but the following code worked for me in a similar task:

import tensorflow as tf
import open3d.ml.tf as ml3d
from open3d.ml.tf.layers import FixedRadiusSearch
from open3d.ml.tf.ops import batch_grid_subsampling as tf_batch_subsampling

nsearch = FixedRadiusSearch()

indx = 10 # indx of point cloud in dataset
p_cloud = point_clouds[indx] # load point cloud

result = nsearch(p_cloud,p_cloud,0.1) # fixed radius search with queries = support points and radius 0.1

idx = result.neighbors_index
splits = result.neighbors_row_splits

ragged_idx = tf.RaggedTensor.from_row_splits(idx, splits) # -> tensor with shape [num_query_points, <varies>]
dense_idx = ragged_idx.to_tensor(default_value=tf.shape(p_cloud)[0]) # -> tensor with shape [num_query_points, num_max_neighbors]

so e.g. for point p_cloud[x] you get the corresponding neighbor indexes with ragged_idx[x] with a variable number of points or dense_idx[x] with a fixed number of points (with appended 'fake' points)

Hope this helps.