CODEWITHSUNDEEP
pythontensorflowmean
Felipe Moser
Felipe Moser

Reputation: 365

Finding centre of mass of tensor (tensorflow)

Is there an efficient way to find the centre of mass of a tensor? I'm working with N stacked volumes (Nx64x64x64) and would like to obtain an Nx3 tensor with the x,y,z position of the centre of mass of each 64x64x64 volume

Upvotes: 3

Views: 2078

Answers (2)

javidcf
javidcf

Reputation: 59731

Following the formula, you should just need to multiply each coordinate by the corresponding mass, sum everything and divide by the total mass:

import tensorflow as tf

# Input volumes
volumes = tf.placeholder(tf.float32, [None, 64, 64, 64])
# Make array of coordinates (each row contains three coordinates)
ii, jj, kk = tf.meshgrid(tf.range(64), tf.range(64), tf.range(64), indexing='ij')
coords = tf.stack([tf.reshape(ii, (-1,)), tf.reshape(jj, (-1,)), tf.reshape(kk, (-1,))], axis=-1)
coords = tf.cast(coords, tf.float32)
# Rearrange input into one vector per volume
volumes_flat = tf.reshape(volumes, [-1, 64 * 64 * 64, 1])
# Compute total mass for each volume
total_mass = tf.reduce_sum(volumes_flat, axis=1)
# Compute centre of mass
centre_of_mass = tf.reduce_sum(volumes_flat * coords, axis=1) / total_mass

Upvotes: 4

Learning is a mess
Learning is a mess

Reputation: 8277

This is a perfect use case for numpy.apply_over_axes:

my_tensor = np.apply_over_axes( np.mean, my_tensor, (1,2,3))

This will return an array of shape (N,1,1,1) with the average along each axis.

Upvotes: 0

Related Questions