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 _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 radon_fft(x):
x_shape = tf.shape(x)
n_angles = x_shape[1]
n_cols = x_shape[2]
x = tf.reshape(x, (-1, n_cols))
x = tf.cast(x, tf.complex64)
x = tf.spectral.fft(x)
x = tf.abs(x)
x = tf.reshape(x, (-1, n_angles, n_cols, 1))
return x