def compute_edges_principled_eager(coords,
normals,
depth=4,
k0=16,
tree_impl=DEFAULT_TREE):
coords = coords.numpy() if hasattr(coords, 'numpy') else coords
normals = normals.numpy() if hasattr(normals, 'numpy') else normals
tree = tree_impl(coords)
dists, indices = tree.query(tree.data, 2, return_distance=True)
# closest = np.min(dists[:, 1])
scale = np.mean(dists[:, 1])
assert (scale > 0)
coords *= (2 / scale)
# coords is now a packing of barely-intersecting spheres of radius 1.
all_coords = [coords]
if normals is not None:
all_normals = [normals]
tree = tree_impl(coords)
trees = [tree]
# base_coords = coords
# base_tree = tree
# base_normals = normals
# perform sampling, build trees
radii = 4 * np.power(2, np.arange(depth))
## Rejection sample on original cloud
# for i, radius in enumerate(radii[:-1]):
# indices = core.rejection_sample_lazy(base_tree,
# base_coords,
# radius,
# k0=k0 * 4**i)
# coords = base_coords[indices]
# tree = tree_impl(coords)
# all_coords.append(coords)
# trees.append(tree)
# if normals is not None:
# all_normals.append(base_normals[indices])
# Recursively rejection sample
# Significantly faster, but not -all- top-level nodes will be in all
# neighborhoords.
for i, radius in enumerate(radii[:-1]):
indices = core.rejection_sample_active(tree, coords, radius, k0=k0)
coords = coords[indices]
tree = tree_impl(coords)
all_coords.append(coords)
trees.append(tree)
if normals is not None:
all_normals.append(normals[indices])
# compute edges
flat_node_indices = utils.lower_triangular(depth)
flat_rel_coords = utils.lower_triangular(depth)
row_splits = utils.lower_triangular(depth)
for i in range(depth):
for j in range(i + 1):
indices = trees[i].query_ball_point(all_coords[j],
radii[i],
approx_neighbors=k0)
flat_node_indices[i][j] = fni = indices.flat_values.astype(
np.int64)
row_splits[i][j] = indices.row_splits.astype(np.int64)
# compute flat_rel_coords
# this could be done outside the py_function, but it uses np.repeat
# which is faster than tf.repeat on cpu.
flat_rel_coords[i][j] = np_cloud.get_relative_coords(
all_coords[i],
all_coords[j],
fni,
row_lengths=indices.row_lengths)
if normals is None:
return (all_coords, flat_rel_coords, flat_node_indices, row_splits)
else:
return (all_coords, all_normals, flat_rel_coords, flat_node_indices,
row_splits)
def post_batch_map(features, labels, weights=None):
all_coords, rel_coords, node_indices = (features[k]
for k in ('all_coords',
'rel_coords',
'node_indices'))
# remove redundant ragged dimension added for batching purposes.
all_coords = tf.nest.map_structure(
lambda x: tf.RaggedTensor.from_nested_row_splits(
x.flat_values, x.nested_row_splits[1:]), all_coords)
offsets = [op_utils.get_row_offsets(c) for c in all_coords]
outer_row_splits = tuple(op_utils.get_row_splits(c) for c in all_coords)
row_splits = tf.nest.map_structure(lambda rt: rt.nested_row_splits[1],
node_indices)
all_coords = tf.nest.map_structure(layer_utils.flatten_leading_dims,
all_coords)
flat_rel_coords = tf.nest.map_structure(lambda x: x.flat_values,
rel_coords)
K = len(all_coords)
flat_node_indices = utils.lower_triangular(K)
for i in range(K):
for j in range(i + 1):
flat_node_indices[i][j] = layer_utils.apply_row_offset(
node_indices[i][j], offsets[i]).flat_values
flat_node_indices = utils.ttuple(flat_node_indices)
all_coords = tuple(all_coords)
class_index = features.get('class_index')
class_masks = features.get('class_masks')
features = dict(all_coords=all_coords,
flat_rel_coords=flat_rel_coords,
flat_node_indices=flat_node_indices,
row_splits=row_splits,
outer_row_splits=outer_row_splits)
all_normals = features.get('all_normals')
if all_normals is not None:
all_normals = tf.nest.map_structure(
lambda x: tf.RaggedTensor.from_nested_row_splits(
x.flat_values, x.nested_row_splits[1:]), all_normals)
all_normals = tf.nest.map_structure(layer_utils.flatten_leading_dims,
all_normals)
all_normals = tuple(all_normals)
features['all_normals'] = all_normals
if class_index is not None:
features['class_index'] = class_index
if class_masks is not None:
features['class_masks'] = class_masks
labels, weights = get_current_problem().post_batch_map(labels, weights)
if isinstance(labels, tf.Tensor) and labels.shape.ndims == 2:
assert (isinstance(weights, tf.Tensor) and weights.shape.ndims == 2)
labels = tf.reshape(labels, (-1, ))
weights = tf.reshape(weights, (-1, ))
return ((features, labels) if weights is None else
(features, labels, weights))
def compute_edges_eager(coords,
normals,
radii,
pooler,
reorder=False,
tree_impl=DEFAULT_TREE,
k0=16):
"""
Recursively sample the input cloud and find edges based on ball searches.
Args:
coords: [N_0, num_dims] numpy array or eager tensor of cloud
coordinates.
normals: [N_0, num_features] numpy array of node features.
radii: [depth] numpy array or eager tensor of radii for ball searches.
pooler: `Pooler` instance.
tree_impl: KDTree implementation.
Returns:
See `compute_edges` return values.
"""
from deep_cloud.ops.np_utils import tree_utils
from deep_cloud.ops.np_utils import cloud as np_cloud
depth = len(radii)
# accomodate eager coords tensor, so can be used with tf.py_functions
if hasattr(coords, 'numpy'):
coords = coords.numpy()
if hasattr(radii, 'numpy'):
radii = radii.numpy()
if hasattr(normals, 'numpy'):
normals = normals.numpy()
if any(isinstance(t, tf.Tensor) for t in (coords, normals, radii)):
assert (tf.executing_eagerly())
if reorder:
indices = np_ordering.iterative_farthest_point_ordering(
coords, coords.shape[0] // 2)
if normals is None:
coords = np_ordering.partial_reorder(indices, coords)
else:
coords, normals = np_ordering.partial_reorder(
indices, coords, normals)
node_indices = utils.lower_triangular(depth)
trees = [None for _ in range(depth)]
all_coords = [coords]
all_normals = [normals]
# TODO: The below uses `depth * (depth + 1) // 2` ball searches
# we can do it with `depth`. When depth is 4 it's not that big a saving...
# do the edges diagonal and build up coordinates
# print('---')
for i, radius in enumerate(radii[:-1]):
tree = trees[i] = tree_impl(coords)
indices = node_indices[i][i] = tree.query_pairs(radius,
approx_neighbors=k0 *
2**i)
# print(radius, tree.n)
# print(np.mean(indices.row_lengths))
coords, normals, indices = pooler(coords, normals, indices)
# node_indices[i + 1][i] = indices
all_coords.append(coords)
all_normals.append(normals)
# final cloud
tree = trees[-1] = tree_impl(coords)
node_indices[-1][-1] = tree.query_pairs(radii[-1],
approx_neighbors=k0 *
2**(depth - 1))
# do below the diagonal, i.e. [i, j], i > j + 1
for i in range(1, depth):
in_tree = trees[i]
radius = radii[i]
for j in range(i):
node_indices[i][j] = trees[j].query_ball_tree(in_tree,
radius,
approx_neighbors=k0)
# TODO: Should be able to do the following in `depth` `rel_coords` calls
# Currently uses `depth * (depth + 1) // 2`. Maybe not efficient?
flat_rel_coords = utils.lower_triangular(depth)
for i in range(depth):
for j in range(i + 1):
indices = node_indices[i][j]
flat_rel_coords[i][j] = np_cloud.get_relative_coords(
all_coords[i],
all_coords[j],
indices.flat_values,
row_lengths=indices.row_lengths)
flat_node_indices = tf.nest.map_structure(
lambda x: x.flat_values.astype(np.int64), node_indices)
row_splits = tf.nest.map_structure(lambda x: x.row_splits.astype(np.int64),
node_indices)
if normals is None:
return (all_coords, flat_rel_coords, flat_node_indices, row_splits)
else:
return (all_coords, all_normals, flat_rel_coords, flat_node_indices,
row_splits)
def compute_edges(coords, normals, depth=4, eager_fn=None):
"""
Graph-mode wrapper for compute_edges implementation.
Recursively sample the input cloud and find edges based on ball searches.
Args:
coords: [N_0, num_dims] numpy array or eager tensor of cloud
coordinates.
normals: [N_0, num_features] numpy array of node features.
eager_fn: function mapping eager versions of coords, normals, depth
to outputs described below.
Returns:
all_coords: [depth] list of numpy coordinate arrays of shape
[N_i, num_dims]
all_normals: [depth] list of numpy node features of shape
[N_i, num_features]
flat_rel_coords: flat values in rel_coords (see below)
flat_node_indices: flat values in node_indices (see below)
row_splits: row splits for rel_coords, node_indices (see below)
The following aren't actually returned but are easier to describe this way.
rel_coords: [depth, <=depth] list of lists of [N_i, k?, num_dims]
floats, where k is the number of neighbors (ragged).
`rel_coords[i][j][p]` (k?, num_dims) gives the relative coordinates
of all points in cloud `i` in the neighborhood of point `p` in cloud
`j`.
node_indices: [depth, <=depth] list of lists of [N_i, k?] ints.
see below. `node_indices[i][j][p] == (r, s, t)` indicates that
node `p` in cloud `j` neighbors nodes `r, s, t` in cloud `i`.
"""
if eager_fn is None:
eager_fn = compute_edges_eager_fn()
if tf.executing_eagerly():
return eager_fn(coords, normals)
if normals is None:
Tout = (
[tf.float32] * depth, # coords
utils.lower_triangular(depth, tf.float32), # flat_rel_coords
utils.lower_triangular(depth, tf.int64), # flat_node_indices
utils.lower_triangular(depth, tf.int64), # row_splits
)
Tout_flat = tf.nest.flatten(Tout)
# out_flat = tf.py_function(fn, [coords, normals], Tout_flat)
out_flat = tf.py_function(
functools.partial(_flatten_output, eager_fn, normals=None),
[coords], Tout_flat)
(all_coords, flat_rel_coords, flat_node_indices,
row_splits) = tf.nest.pack_sequence_as(Tout, out_flat)
all_normals = None
else:
Tout = (
[tf.float32] * depth, # coords
[tf.float32] * depth, # normals
utils.lower_triangular(depth, tf.float32), # flat_rel_coords
utils.lower_triangular(depth, tf.int64), # flat_node_indices
utils.lower_triangular(depth, tf.int64), # row_splits
)
Tout_flat = tf.nest.flatten(Tout)
# out_flat = tf.py_function(fn, [coords, normals], Tout_flat)
out_flat = tf.py_function(functools.partial(_flatten_output, eager_fn),
[coords, normals], Tout_flat)
(all_coords, all_normals, flat_rel_coords, flat_node_indices,
row_splits) = tf.nest.pack_sequence_as(Tout, out_flat)
# fix sizes
num_dims = coords.shape[1]
# num_features = normals.shape[1]
# sizes = [coords.shape[0]]
# even if size is None, we still get rank information from the below
# pooler = poolers.SlicePooler()
# for _ in range(depth - 1):
# sizes.append(pooler.output_size(sizes[-1]))
sizes = [None] * depth
for i in range(depth):
size = sizes[i]
all_coords[i].set_shape((size, num_dims))
if normals is not None:
all_normals[i].set_shape((size, normals.shape[1]))
for rc in flat_rel_coords[i]:
rc.set_shape((None, num_dims))
for ni in flat_node_indices[i]:
ni.set_shape((None, ))
for rs in row_splits[i]:
rs.set_shape((None if size is None else (size + 1), ))
return (all_coords, all_normals, flat_rel_coords, flat_node_indices,
row_splits)
def __call__(
self,
node_features,
edge_features,
global_features,
global_edge_features=None,
):
# input validation
for i, efs in enumerate(edge_features):
if not isinstance(efs, (list, tuple)):
raise ValueError(
'entries of edge_features must be lists/tuples, '
'but element {} is {}'.format(i, efs))
if len(efs) != i + 1:
raise ValueError(
'entry {} of bipartite_extractors should have length {} '
'but has length {}'.format(i, i + 1, len(efs)))
for j, ef in enumerate(efs):
assert_flat_tensor('edge_features[{}][{}]'.format(i, j), ef, 2,
FLOAT_TYPES)
extractors = self.bipartite_extractors
K = len(extractors)
if node_features is None:
node_features = [None] * K
elif len(node_features) != K:
raise ValueError('Expected {} node features, got {}'.format(
K, len(node_features)))
for i, nf in enumerate(node_features):
if nf is not None:
assert_flat_tensor('node_features[{}]'.format(i), nf, 2,
FLOAT_TYPES)
# -------------------------
# finished validating inputs
# -------------------------
all_out_edge_features = utils.lower_triangular(K)
all_out_node_features = [[None] * K for _ in range(K)]
for i in range(K):
for j in range(i):
if extractors[i][j] is None:
continue
af, bf, ef = extractors[i][j](node_features[j],
node_features[i],
edge_features[i][j])
all_out_node_features[i][j] = bf
all_out_node_features[j][i] = af
all_out_edge_features[i][j] = ef
# symmetric cloud
af, bf, ef = extractors[i][i](node_features[i],
None,
edge_features[i][i],
symmetric=True)
assert (bf is None)
all_out_node_features[i][i] = af
all_out_edge_features[i][i] = ef
# add global features
if self.global_extractors is not None:
for i, extractor in enumerate(self.global_extractors):
if extractor is not None:
ef = (None if global_edge_features is None else
global_edge_features[i])
all_out_node_features[i].append(
extractor(global_features, node_features[i], ef))
# concatenate node features
for i in range(K):
all_out_node_features[i] = tf.concat(all_out_node_features[i],
axis=-1)
return all_out_node_features, all_out_edge_features