def test_fixed_radius_search(dtype, ml, num_points_queries, radius,
hash_table_size_factor, metric,
ignore_query_point, return_distances):
# skip dtype not supported on GPU
if mltest.is_gpu_device_name(ml.device) and not dtype in gpu_dtypes:
return
rng = np.random.RandomState(123)
num_points, num_queries = num_points_queries
points = rng.random(size=(num_points, 3)).astype(dtype)
if ignore_query_point:
queries = points
else:
queries = rng.random(size=(num_queries, 3)).astype(dtype)
# kd tree for computing the ground truth
tree = cKDTree(points, copy_data=True)
p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric]
gt_neighbors_index = tree.query_ball_point(queries, radius, p=p_norm)
layer = ml.layers.FixedRadiusSearch(metric=metric,
ignore_query_point=ignore_query_point,
return_distances=return_distances)
ans = mltest.run_op(
ml,
ml.device,
True,
layer,
points,
queries=queries,
radius=radius,
hash_table_size_factor=hash_table_size_factor,
)
for i, q in enumerate(queries):
# check neighbors
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
q_neighbors_index = ans.neighbors_index[start:end]
gt_set = set(gt_neighbors_index[i])
if ignore_query_point:
gt_set.remove(i)
assert gt_set == set(q_neighbors_index)
# check distances
if return_distances:
q_neighbors_dist = ans.neighbors_distance[start:end]
for j, dist in zip(q_neighbors_index, q_neighbors_dist):
if metric == 'L2':
gt_dist = np.sum((q - points[j])**2)
else:
gt_dist = np.linalg.norm(q - points[j], ord=p_norm)
np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
# ----------------------------------------------------------------------------
import open3d as o3d
import numpy as np
import pytest
import mltest
# Skip all tests if the ml ops were not built.
pytestmark = mltest.default_marks
ml_torch_gpu_only = pytest.mark.parametrize('ml', [
v for k, v in mltest._ml_modules.items()
if mltest.is_gpu_device_name(v.device) and v.module.__name__ == 'torch'
])
@ml_torch_gpu_only
def test_group_pts(ml):
values0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/values0.npy'
)
values1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/values1.npy'
)
ans = mltest.run_op(ml, ml.device, True, ml.ops.group_points, values0,
values1)
def test_fixed_radius_search_raggedtensor(dtype, ml, batch_size, radius,
hash_table_size_factor, metric,
ignore_query_point, return_distances,
index_dtype):
# the problem is specific to tensorflow
if ml.module.__name__ != 'tensorflow':
return
# skip dtype not supported on GPU
if mltest.is_gpu_device_name(ml.device) and not dtype in gpu_dtypes:
return
rng = np.random.RandomState(123)
# create array defining start and end of each batch
points_row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64)
queries_row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64)
for i in range(batch_size):
points_row_splits[i + 1] = rng.randint(15) + points_row_splits[i]
queries_row_splits[i + 1] = rng.randint(15) + queries_row_splits[i]
num_points = points_row_splits[-1]
num_queries = queries_row_splits[-1]
points = rng.random(size=(num_points, 3)).astype(dtype)
if ignore_query_point:
queries = points
queries_row_splits = points_row_splits
else:
queries = rng.random(size=(num_queries, 3)).astype(dtype)
# kd trees for computing the ground truth
p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric]
gt_neighbors_index = []
for i in range(batch_size):
points_i = points[points_row_splits[i]:points_row_splits[i + 1]]
queries_i = queries[queries_row_splits[i]:queries_row_splits[i + 1]]
tree = cKDTree(points_i, copy_data=True)
gt_neighbors_index.extend([
list(
tree.query_ball_point(q, radius, p=p_norm) +
points_row_splits[i]) for q in queries_i
])
if ml.module.__name__ == 'tensorflow':
index_dtype_ = {'int32': tf.int32, 'int64': tf.int64}[index_dtype]
elif ml.module.__name__ == 'torch':
index_dtype_ = {
'int32': torch.int32,
'int64': torch.int64
}[index_dtype]
else:
raise Exception('Unsupported ml framework')
points_ragged = tf.RaggedTensor.from_row_splits(
values=points, row_splits=points_row_splits)
queries_ragged = tf.RaggedTensor.from_row_splits(
values=queries, row_splits=queries_row_splits)
layer = ml.layers.FixedRadiusSearch(metric=metric,
ignore_query_point=ignore_query_point,
return_distances=return_distances,
index_dtype=index_dtype_)
ans = mltest.run_op(
ml,
ml.device,
True,
layer,
points_ragged,
queries=queries_ragged,
radius=radius,
hash_table_size_factor=hash_table_size_factor,
)
index_dtype_np = {'int32': np.int32, 'int64': np.int64}[index_dtype]
assert ans.neighbors_index.dtype == index_dtype_np
for i, q in enumerate(queries):
# check neighbors
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
q_neighbors_index = ans.neighbors_index[start:end]
gt_set = set(gt_neighbors_index[i])
if ignore_query_point:
gt_set.remove(i)
assert gt_set == set(q_neighbors_index)
# check distances
if return_distances:
q_neighbors_dist = ans.neighbors_distance[start:end]
for j, dist in zip(q_neighbors_index, q_neighbors_dist):
if metric == 'L2':
gt_dist = np.sum((q - points[j])**2)
else:
gt_dist = np.linalg.norm(q - points[j], ord=p_norm)
np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)