def tf_map_coordinates(input, coords, order=1):
"""Tensorflow verion of scipy.ndimage.map_coordinates"""
assert order == 1
coords_tl = tf.cast(tf.floor(coords), 'int32')
coords_br = tf.cast(tf.ceil(coords), 'int32')
coords_bl = tf.stack([coords_tl[:, 0], coords_br[:, 1]], axis=1)
coords_tr = tf.stack([coords_br[:, 0], coords_tl[:, 1]], axis=1)
vals_tl = tf.gather_nd(input, coords_tl)
vals_br = tf.gather_nd(input, coords_br)
vals_bl = tf.gather_nd(input, coords_bl)
vals_tr = tf.gather_nd(input, coords_tr)
h_offset = coords[:, 0] - tf.cast(coords_tl[:, 0], tf.float32)
h_int_t = (((1.0 - h_offset) * vals_tl) + (h_offset * vals_tr))
h_int_b = (((1.0 - h_offset) * vals_bl) + (h_offset * vals_br))
v_offset = coords[:, 1] - tf.cast(coords_tl[:, 1], tf.float32)
int_vals = (((1.0 - v_offset) * h_int_t) + (v_offset * h_int_b))
return int_vals
def top_k(scores, I, ratio, top_k_var):
"""
Returns indices to get the top K values in `scores` segment-wise, with
segments defined by I. K is not fixed, but it is defined as a ratio of the
number of elements in each segment.
:param scores: a rank 1 tensor with scores;
:param I: a rank 1 tensor with segment IDs;
:param ratio: float, ratio of elements to keep for each segment;
:param top_k_var: a tf.Variable without shape validation (e.g.,
`tf.Variable(0.0, validate_shape=False)`);
:return: a rank 1 tensor containing the indices to get the top K values of
each segment in `scores`.
"""
num_nodes = tf.segment_sum(tf.ones_like(I),
I) # Number of nodes in each graph
cumsum = tf.cumsum(num_nodes) # Cumulative number of nodes (A, A+B, A+B+C)
cumsum_start = cumsum - num_nodes # Start index of each graph
n_graphs = tf.shape(num_nodes)[0] # Number of graphs in batch
max_n_nodes = tf.reduce_max(num_nodes) # Order of biggest graph in batch
batch_n_nodes = tf.shape(I)[0] # Number of overall nodes in batch
to_keep = tf.ceil(ratio * tf.cast(num_nodes, tf.float32))
to_keep = tf.cast(to_keep, tf.int32) # Nodes to keep in each graph
index = tf.range(batch_n_nodes)
index = (index - tf.gather(cumsum_start, I)) + (I * max_n_nodes)
y_min = tf.reduce_min(scores)
dense_y = tf.ones((n_graphs * max_n_nodes, ))
dense_y = dense_y * tf.cast(
y_min - 1, tf.float32
) # subtract 1 to ensure that filler values do not get picked
dense_y = tf.assign(
top_k_var, dense_y, validate_shape=False
) # top_k_var is a variable with unknown shape defined in the elsewhere
dense_y = tf.scatter_update(dense_y, index, scores)
dense_y = tf.reshape(dense_y, (n_graphs, max_n_nodes))
perm = tf.argsort(dense_y, direction='DESCENDING')
perm = perm + cumsum_start[:, None]
perm = tf.reshape(perm, (-1, ))
to_rep = tf.tile(tf.constant([1., 0.]), (n_graphs, ))
rep_times = tf.reshape(
tf.concat((to_keep[:, None], (max_n_nodes - to_keep)[:, None]), -1),
(-1, ))
mask = tf_repeat_1d(to_rep, rep_times)
perm = tf.boolean_mask(perm, mask)
return perm
def batch_map_coordinates(input, coords, order=1):
"""Batch version of tf_map_coordinates"""
input_shape = tf.shape(input)
batch_size = input_shape[0]
input_size = input_shape[1]
#coords = tf.reshape(coords, (batch_size, -1, 2))
n_coords = tf.shape(coords)[1]
coords = tf.clip_by_value(coords, 0, tf.cast(input_size, 'float32') - 1)
coords_tl = tf.cast(tf.floor(coords), 'int32')
coords_br = tf.cast(tf.ceil(coords), 'int32')
coords_bl = tf.stack([coords_tl[..., 0], coords_br[..., 1]], axis=-1)
coords_tr = tf.stack([coords_br[..., 0], coords_tl[..., 1]], axis=-1)
idx = tf.range(batch_size)
idx = tf.expand_dims(idx, -1)
idx = tf.tile(idx, [1, n_coords])
idx = tf.reshape(idx, [-1])
def _get_vals_by_coords(input, coords):
coords_0_flat = tf.reshape(coords[..., 0], [-1])
coords_1_flat = tf.reshape(coords[..., 1], [-1])
indices = tf.stack([idx, coords_0_flat, coords_1_flat], axis=-1)
vals = tf.gather_nd(input, indices)
vals = tf.reshape(vals, (batch_size, n_coords))
return vals
vals_tl = _get_vals_by_coords(input, coords_tl)
vals_br = _get_vals_by_coords(input, coords_br)
vals_bl = _get_vals_by_coords(input, coords_bl)
vals_tr = _get_vals_by_coords(input, coords_tr)
h_offset = coords[..., 0] - tf.cast(coords_tl[..., 0], tf.float32)
h_int_t = (((1.0 - h_offset) * vals_tl) + (h_offset * vals_tr))
h_int_b = (((1.0 - h_offset) * vals_bl) + (h_offset * vals_br))
v_offset = coords[..., 1] - tf.cast(coords_tl[..., 1], tf.float32)
int_vals = (((1.0 - v_offset) * h_int_t) + (v_offset * h_int_b))
return int_vals
def _upsampled_registration(target_image, src_image, upsample_factor):
upsample_factor = tf.constant(upsample_factor, tf.float32)
target_shape = tf.shape(target_image)
target_image = tf.reshape(target_image, target_shape[:3])
src_shape = tf.shape(src_image)
src_image = tf.reshape(src_image, src_shape[:3])
src_freq = fft2d(src_image)
target_freq = fft2d(target_image)
shape = tf.reshape(tf.shape(src_freq)[1:3], (1, 2))
shape = tf.cast(shape, tf.float32)
shape = tf.tile(shape, (tf.shape(target_freq)[0], 1))
image_product = src_freq * tf.conj(target_freq)
cross_correlation = tf.spectral.ifft2d(image_product)
maxima = find_maxima(tf.abs(cross_correlation))
midpoints = fix(tf.cast(shape, tf.float32) / 2.)
shifts = maxima
shifts = tf.where(shifts > midpoints, shifts - shape, shifts)
shifts = tf.round(shifts * upsample_factor) / upsample_factor
upsampled_region_size = tf.ceil(upsample_factor * 1.5)
dftshift = fix(upsampled_region_size / 2.0)
normalization = tf.cast(tf.size(src_freq[0]), tf.float32)
normalization *= upsample_factor**2
sample_region_offset = dftshift - shifts * upsample_factor
data = tf.conj(image_product)
upsampled_dft = _upsampled_dft(data, upsampled_region_size,
upsample_factor, sample_region_offset)
cross_correlation = tf.conj(upsampled_dft)
cross_correlation /= tf.cast(normalization, tf.complex64)
cross_correlation = tf.abs(cross_correlation)
maxima = find_maxima(cross_correlation)
maxima = maxima - dftshift
shifts = shifts + maxima / upsample_factor
return shifts
def batch_map_coordinates(input, coords, n_coords):
"""Batch version of tf_map_coordinates"""
#init_input_shape = input.get_shape()
#input = tf.reshape(input, (-1, init_input_shape[3]))
#coords = tf.reshape(coords, (-1, n_coords * 2))
input_shape = input.get_shape()
input_h = input_shape[1].value
input_w = input_shape[2].value
#batch_size = input_shape[0]
#input_size = input_shape[1]
#coords = tf.reshape(coords, (batch_size, -1, 2))
#n_coords = tf.shape(coords)[1]
coords_h = tf.clip_by_value(coords[..., :n_coords], 0,
tf.cast(input_h, 'float32') - 1)
coords_w = tf.clip_by_value(coords[..., n_coords:], 0,
tf.cast(input_w, 'float32') - 1)
coords = tf.stack([coords_h, coords_w], axis=-1)
coords_tl = tf.cast(tf.floor(coords), 'float32')
coords_br = tf.cast(tf.ceil(coords), 'float32')
coords_bl = tf.stack([coords_tl[..., 0], coords_br[..., 1]], axis=-1)
coords_tr = tf.stack([coords_br[..., 0], coords_tl[..., 1]], axis=-1)
#idx = tf.range(batch_size)
#idx = tf.expand_dims(idx, -1)
#idx = tf.tile(idx, [1, n_coords])
#idx = tf.reshape(idx, [-1])
def _get_vals_by_coords(input, coords, n_coords):
coords_shape = tf.shape(coords)
input_shape = input.get_shape()
input_w = input_shape[2].value
input_h = input_shape[1].value
channel_size = input_shape[3].value
batch_size = tf.shape(input)[0]
input = tf.transpose(input, (0, 3, 1, 2))
input = tf.reshape(input, (-1, channel_size, input_h * input_w))
indices = coords[..., 0] * input_w + coords[..., 1]
#indices = tf.expand_dims(indices, axis=1)
#indices = tf.tile(indices, [1, channel_size, 1, 1, 1])
#indices = tf.reshape(indices, (-1, channel_size, input_h * input_w * n_coords))
#indices = tf.transpose(indices, (0, 3, 1, 2))
indices = tf.reshape(indices, (-1, input_h * input_w * n_coords))
indices = tf.cast(indices, 'int32')
#indices = tf.reshape(indices, [-1])
#input = tf.reshape(input, [-1])
vals = tf.gather(input, indices[0], axis=-1)
#vals = tf.map_fn(lambda x: tf.gather(x[0], x[1], axis=-1), (input,indices), dtype=tf.float32)
vals = tf.reshape(vals, (-1, channel_size, input_h, input_w, n_coords))
return vals
vals_tl = (1 + (coords_tl[..., 0] - coords[..., 0])) * \
(1 + (coords_tl[..., 1] - coords[..., 1]))
vals_br = (1 - (coords_br[..., 0] - coords[..., 0])) * \
(1 - (coords_br[..., 1] - coords[..., 1]))
vals_bl = (1 + (coords_bl[..., 0] - coords[..., 0])) * \
(1 + (coords_bl[..., 1] - coords[..., 1]))
vals_tr = (1 - (coords_tr[..., 0] - coords[..., 0])) * \
(1 - (coords_tr[..., 1] - coords[..., 1]))
x_vals_tl = _get_vals_by_coords(input, coords_tl, n_coords)
x_vals_br = _get_vals_by_coords(input, coords_br, n_coords)
x_vals_bl = _get_vals_by_coords(input, coords_bl, n_coords)
x_vals_tr = _get_vals_by_coords(input, coords_tr, n_coords)
#h_offset = coords[..., 0] - tf.cast(coords_tl[..., 0], tf.float32)
#h_int_t = (((1.0 - h_offset) * vals_tl) + (h_offset * vals_tr))
#h_int_b = (((1.0 - h_offset) * vals_bl) + (h_offset * vals_br))
#v_offset = coords[..., 1] - tf.cast(coords_tl[..., 1], tf.float32)
#int_vals = (((1.0 - v_offset) * h_int_t) + (v_offset * h_int_b))
int_vals = tf.expand_dims(vals_tl, 1) * x_vals_tl + \
tf.expand_dims(vals_br, 1) * x_vals_br + \
tf.expand_dims(vals_bl, 1) * x_vals_bl + \
tf.expand_dims(vals_tr, 1) * x_vals_tr
return int_vals
def fix(x):
x = tf.where(x >= 0, tf.floor(x), tf.ceil(x))
return x
