def _find_subpixel_maxima(x,
kernel_size,
sigma,
upsample_factor,
coordinate_scale=1,
confidence_scale=255.):
kernel = gaussian_kernel_2d(kernel_size, sigma)
kernel = tf.expand_dims(kernel, 0)
x_shape = tf.shape(x)
rows = x_shape[1]
cols = x_shape[2]
max_vals = tf.reduce_max(tf.reshape(x, [-1, rows * cols]), axis=1)
max_vals = tf.reshape(max_vals, [-1, 1]) / confidence_scale
row_pad = rows // 2 - kernel_size // 2
col_pad = cols // 2 - kernel_size // 2
padding = [[0, 0], [row_pad, row_pad - 1], [col_pad, col_pad - 1]]
kernel = tf.pad(kernel, padding)
row_center = row_pad + (kernel_size // 2)
col_center = col_pad + (kernel_size // 2)
center = tf.stack([row_center, col_center])
center = tf.expand_dims(center, 0)
center = tf.cast(center, dtype=tf.float32)
shifts = _upsampled_registration(x, kernel, upsample_factor)
shifts = center - shifts
shifts *= coordinate_scale
maxima = tf.concat([shifts[:, ::-1], max_vals], -1)
return maxima
def sparse_bool_mask(x, mask, axis=0):
# Only necessary if indices may have non-unique elements
indices = tf.boolean_mask(tf.range(tf.shape(x)[axis]), mask)
n_indices = tf.size(indices)
# Get indices for the axis
idx = x.indices[:, axis]
# Find where indices match the selection
eq = tf.equal(tf.expand_dims(idx, 1),
tf.cast(indices, tf.int64)) # TODO this has quadratic cost
# Mask for selected values
sel = tf.reduce_any(eq, axis=1)
# Selected values
values_new = tf.boolean_mask(x.values, sel, axis=0)
# New index value for selected elements
n_indices = tf.cast(n_indices, tf.int64)
idx_new = tf.reduce_sum(tf.cast(eq, tf.int64) * tf.range(n_indices),
axis=1)
idx_new = tf.boolean_mask(idx_new, sel, axis=0)
# New full indices tensor
indices_new = tf.boolean_mask(x.indices, sel, axis=0)
indices_new = tf.concat([
indices_new[:, :axis],
tf.expand_dims(idx_new, 1), indices_new[:, axis + 1:]
],
axis=1)
# New shape
shape_new = tf.concat(
[x.dense_shape[:axis], [n_indices], x.dense_shape[axis + 1:]], axis=0)
return tf.SparseTensor(indices_new, values_new, shape_new)
def call(self, inputs):
print("xxxx",inputs)
expanded_tensor = ktf.expand_dims(inputs[0], -1)
multiples = [1, self.number_of_transforms, 1, 1, 1]
tiled_tensor = ktf.tile(expanded_tensor, multiples=multiples)
repeated_tensor = ktf.reshape(tiled_tensor, ktf.shape(inputs[0]) * np.array([self.number_of_transforms, 1, 1, 1]))
affine_transforms = inputs[1] / self.affine_mul
affine_transforms = ktf.reshape(affine_transforms, (-1, 8))
tranformed = tf_affine_transform(repeated_tensor, affine_transforms)
res = ktf.reshape(tranformed, [-1, self.number_of_transforms] + self.image_size)
res = ktf.transpose(res, [0, 2, 3, 1, 4])
#Use masks
if len(inputs) == 3:
mask = ktf.transpose(inputs[2], [0, 2, 3, 1])
mask = ktf.image.resize_images(mask, self.image_size[:2], method=ktf.image.ResizeMethod.NEAREST_NEIGHBOR)
res = res * ktf.expand_dims(mask, axis=-1)
if self.aggregation_fn == 'none':
res = ktf.reshape(res, [-1] + self.image_size[:2] + [self.image_size[2] * self.number_of_transforms])
elif self.aggregation_fn == 'max':
res = ktf.reduce_max(res, reduction_indices=[-2])
elif self.aggregation_fn == 'avg':
counts = ktf.reduce_sum(mask, reduction_indices=[-1])
counts = ktf.expand_dims(counts, axis=-1)
res = ktf.reduce_sum(res, reduction_indices=[-2])
res /= counts
res = ktf.where(ktf.is_nan(res), ktf.zeros_like(res), res)
return res
def _col_kernel(upsampled_region_size, upsample_factor, axis_offsets,
data_shape):
data_shape_float = tf.cast(data_shape, tf.float32)
col_constant = tf.cast(data_shape_float[2] * upsample_factor, tf.complex64)
col_constant = (-1j * 2 * np.pi / col_constant)
col_kernel_a = tf.range(0, data_shape_float[2], dtype=tf.float32)
col_kernel_a = fftshift1d(col_kernel_a)
col_kernel_a = tf.reshape(col_kernel_a, (-1, 1))
col_kernel_a -= tf.floor(data_shape_float[2] / 2.)
col_kernel_a = tf.reshape(col_kernel_a, (1, -1))
col_kernel_a = tf.tile(col_kernel_a, (data_shape[0], 1))
col_kernel_b = tf.range(0, upsampled_region_size, dtype=tf.float32)
col_kernel_b = tf.reshape(col_kernel_b, (1, -1))
col_kernel_b = tf.tile(col_kernel_b, (data_shape[0], 1))
col_kernel_b = tf.transpose(col_kernel_b)
col_kernel_b -= tf.transpose(axis_offsets[:, 1])
col_kernel_b = tf.transpose(col_kernel_b)
col_kernel_a = tf.expand_dims(col_kernel_a, 1)
col_kernel_b = tf.expand_dims(col_kernel_b, -1)
col_kernel = col_kernel_a * col_kernel_b
col_kernel = tf.transpose(col_kernel, perm=(0, 2, 1))
col_kernel = col_constant * tf.cast(col_kernel, tf.complex64)
col_kernel = tf.exp(col_kernel)
return col_kernel
def _row_kernel(upsampled_region_size, upsample_factor, axis_offsets,
data_shape):
data_shape_float = tf.cast(data_shape, tf.float32)
row_constant = tf.cast(data_shape_float[1] * upsample_factor, tf.complex64)
row_constant = (-1j * 2 * np.pi / row_constant)
row_kernel_a = tf.range(0, upsampled_region_size, dtype=tf.float32)
row_kernel_a = tf.reshape(row_kernel_a, (1, -1))
row_kernel_a = tf.tile(row_kernel_a, (data_shape[0], 1))
row_kernel_a = tf.transpose(row_kernel_a)
row_kernel_a = row_kernel_a - axis_offsets[:, 0]
row_kernel_b = tf.range(0, data_shape_float[1], dtype=tf.float32)
row_kernel_b = fftshift1d(row_kernel_b)
row_kernel_b = tf.reshape(row_kernel_b, (1, -1))
row_kernel_b = tf.tile(row_kernel_b, (data_shape[0], 1))
row_kernel_b = row_kernel_b - tf.floor(data_shape_float[1] / 2.)
row_kernel_a = tf.expand_dims(row_kernel_a, 1)
row_kernel_b = tf.expand_dims(row_kernel_b, -1)
row_kernel = tf.transpose(row_kernel_a) * row_kernel_b
row_kernel = tf.transpose(row_kernel, perm=(0, 2, 1))
row_kernel = row_constant * tf.cast(row_kernel, tf.complex64)
row_kernel = tf.exp(row_kernel)
return row_kernel
def nn_loss(self, reference, target, neighborhood_size=(3, 3)):
v_pad = neighborhood_size[0] // 2
h_pad = neighborhood_size[1] // 2
val_pad = ktf.pad(reference,
[[0, 0], [v_pad, v_pad], [h_pad, h_pad], [0, 0]],
mode='CONSTANT',
constant_values=-10000)
reference_tensors = []
for i_begin in range(0, neighborhood_size[0]):
i_end = i_begin - neighborhood_size[0] + 1
i_end = None if i_end == 0 else i_end
for j_begin in range(0, neighborhood_size[1]):
j_end = j_begin - neighborhood_size[0] + 1
j_end = None if j_end == 0 else j_end
sub_tensor = val_pad[:, i_begin:i_end, j_begin:j_end, :]
reference_tensors.append(ktf.expand_dims(sub_tensor, -1))
reference = ktf.concat(reference_tensors, axis=-1)
target = ktf.expand_dims(target, axis=-1)
abs = ktf.abs(reference - target)
norms = ktf.reduce_sum(abs, reduction_indices=[-2])
loss = ktf.reduce_min(norms, reduction_indices=[-1])
return loss
def call(self, inputs):
mask = ktf.transpose(inputs[2], [0, 2, 3, 1])
# mask = ktf.image.resize_images(mask, self.image_size[:2], method=ktf.image.ResizeMethod.NEAREST_NEIGHBOR)
if self.number_of_transforms == 11:
# return inputs[0] * K.tile(ktf.expand_dims(1-mask[:,:,:,0], axis=3),(1,1,1,inputs[0].shape[-1]))
return inputs[0] * ktf.expand_dims(1 - mask[:, :, :, 0], axis=3)
else:
return inputs[0] * ktf.expand_dims(
K.clip(1 - K.sum(mask, 3), 0, 1), axis=3)
def tf_repeat_1d(x, repeats):
"""
Repeats each value `x[i]` a number of times `repeats[i]`.
:param x: a rank 1 tensor;
:param repeats: a rank 1 tensor;
:return: a rank 1 tensor, of shape `(sum(repeats), )`.
"""
x = tf.expand_dims(x, 1)
max_repeats = tf.reduce_max(repeats)
tile_repeats = [1, max_repeats]
arr_tiled = tf.tile(x, tile_repeats)
mask = tf.less(tf.range(max_repeats), tf.expand_dims(repeats, 1))
result = tf.reshape(tf.boolean_mask(arr_tiled, mask), [-1])
return result
def gaussian_kernel_1d(size, sigma):
size = tf.constant(size, dtype=tf.float32)
sigma = tf.constant(sigma, dtype=tf.float32)
x = tf.range(-(size // 2), (size // 2) + 1, dtype=tf.float32)
kernel = 1 / (sigma * tf.sqrt(2 * np.pi))
kernel *= tf.exp(-0.5 * (x / sigma)**2)
return tf.expand_dims(kernel, axis=-1)
def call(self, inputs):
image = inputs[0]
if len(inputs) == 2:
style = inputs[1]
style_mean, style_var = ktf.nn.moments(style,
self.spatial_axis,
keep_dims=True)
else:
style_mean = ktf.expand_dims(
ktf.expand_dims(inputs[1], self.spatial_axis[0]),
self.spatial_axis[1])
style_var = ktf.expand_dims(
ktf.expand_dims(inputs[2], self.spatial_axis[0]),
self.spatial_axis[1])
image_mean, image_var = ktf.nn.moments(image,
self.spatial_axis,
keep_dims=True)
out = ktf.nn.batch_normalization(image, image_mean,
image_var, style_mean,
ktf.sqrt(style_var), self.eps)
return out
def call(self, inputs):
if len(inputs) == 3:
X, A, I = inputs
self.data_mode = 'graph'
else:
X, A = inputs
I = tf.zeros(tf.shape(X)[:1], dtype=tf.int32)
self.data_mode = 'single'
if K.ndim(I) == 2:
I = I[:, 0]
A_is_sparse = K.is_sparse(A)
# Get mask
y = K.dot(X, K.l2_normalize(self.kernel))
N = K.shape(X)[-2]
indices = ops.top_k(y[:, 0], I, self.ratio, self.top_k_var)
mask = tf.scatter_nd(tf.expand_dims(indices, 1), tf.ones_like(indices),
(N, ))
# Multiply X and y to make layer differentiable
features = X * self.gating_op(y)
axis = 0 if len(K.int_shape(
A)) == 2 else 1 # Cannot use negative axis in tf.boolean_mask
# Reduce X
X_pooled = tf.boolean_mask(features, mask, axis=axis)
# Compute A^2
if A_is_sparse:
A_dense = tf.sparse.to_dense(A)
else:
A_dense = A
A_squared = K.dot(A, A_dense)
# Reduce A
A_pooled = tf.boolean_mask(A_squared, mask, axis=axis)
A_pooled = tf.boolean_mask(A_pooled, mask, axis=axis + 1)
if A_is_sparse:
A_pooled = tf.contrib.layers.dense_to_sparse(A_pooled)
output = [X_pooled, A_pooled]
# Reduce I
if self.data_mode == 'graph':
I_pooled = tf.boolean_mask(I[:, None], mask)[:, 0]
output.append(I_pooled)
if self.return_mask:
output.append(mask)
return output
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 batch_map_offsets(input, offsets, order=1):
"""Batch map offsets into input
Adds index of every entry to the entry to make it's interpolation
relevant to it's location
"""
input_shape = tf.shape(input)
batch_size = input_shape[0]
input_w = input_shape[1]
input_h = input_shape[2]
offsets = tf.reshape(offsets, (batch_size, -1, 2))
ind_add = tf.meshgrid(tf.range(input_w), tf.range(input_h), indexing='ij')
ind_add = tf.stack(ind_add, axis=-1)
ind_add = tf.cast(ind_add, 'float32')
ind_add = tf.reshape(ind_add, (-1, 2))
ind_add = tf.expand_dims(ind_add, 0)
ind_add = tf.tile(ind_add, [batch_size, 1, 1])
coords = offsets + ind_add
int_vals = batch_map_coordinates(input, coords)
return int_vals
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 build(self, input_shape):
self.xx, self.yy = ktf.meshgrid(ktf.range(self.image_size[1]),
ktf.range(self.image_size[0]))
self.xx = ktf.expand_dims(ktf.cast(self.xx, 'float32'), 2)
self.yy = ktf.expand_dims(ktf.cast(self.yy, 'float32'), 2)
