Subtract means collected from tf.math.unsorted_segment_mean from a tensor in Tensorflow

Question

So I have a 2D tf.float32 Tensor of xyz coords and a 1D tf.int32 Tensor of segment_ids. I want to subtract every point from the mean of the corresponding segment_id.

Please check the code below:

x_index = tf.constant([1, 1, 2, 2])
y_index = tf.constant([1, 1, 3, 4])
points = tf.constant([[0.1, 0.1, 0.1], 
                      [0.11, 0.11, 0.11], 
                      [0.2, 0.3, 0.1], 
                      [0.2, 0.4, 0.1]])

points_x_y_indices = tf.transpose(tf.stack([x_index, y_index]))
uniques, idx = tf.raw_ops.UniqueV2(x=points_x_y_indices, axis=[0], out_idx=tf.dtypes.int32)

n_pillars = int(tf.reduce_max(idx))+1
x_means = tf.math.unsorted_segment_mean(points[:, 0], idx, n_pillars)
y_means = tf.math.unsorted_segment_mean(points[:, 1], idx, n_pillars)
z_means = tf.math.unsorted_segment_mean(points[:, 2], idx, n_pillars)

Now, I have the means over every segment_id in x_means, y_means and z_means. How can I subtract those values from original points tensor?? of course without looping as I am trying to avoid tf.py_func

Thanks!

Answer 1

I figured it out, you can use

full_x_means = tf.gather(x_means, idx)
full_y_means = tf.gather(y_means, idx)
full_z_means = tf.gather(z_means, idx)

Then

pillar_points_xc = points[:, 0] - full_x_means
pillar_points_yc = points[:, 1] - full_y_means
pillar_points_zc = points[:, 2] - full_z_means