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