def _register_rotation(target_image, src_image, rotation_resolution,
rotation_guess, upsample_factor):
n_angles = tf.cast(tf.round(180. / rotation_resolution), tf.int32)
theta = tf.linspace(0., 180. - rotation_resolution, n_angles)
theta = -radians(theta)
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])
rotation_guess = tf.constant(rotation_guess, tf.float32)
rotation_resolution = tf.constant(rotation_resolution, tf.float32)
src_image = radon_transform_fft(src_image, theta)
target_image = radon_transform_fft(target_image, theta)
shifts = _upsampled_registration(target_image, src_image, upsample_factor)
angles = shifts[:, 0] * rotation_resolution
angles = tf.reshape(angles, [-1, 1])
angles = check_angles(angles, rotation_guess)
angles = radians(angles)
return angles
def fftshift1d(x, axis=0):
x_shape = tf.shape(x)
x = tf.reshape(x, (-1, 1))
n_samples = tf.cast(tf.shape(x)[0], tf.float32)
even = n_samples / 2.
even = tf.round(even)
even = even * 2.
even = tf.equal(n_samples, even)
def true_fn():
return x
def false_fn():
x_padded = tf.concat([x, tf.zeros((1, 1))], axis=0)
return x_padded
x = tf.cond(even, true_fn, false_fn)
x1, x2 = tf.split(x, 2, axis=axis)
def true_fn():
return x2
def false_fn():
x2_unpadded = x2[:-1]
return x2_unpadded
x2 = tf.cond(even, true_fn, false_fn)
x = tf.concat((x2, x1), axis=axis)
x = tf.reshape(x, x_shape)
return x
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