def hybrid_correlation_tf(img1_fft, img2_fft):
"""Unsure if this actually the hybrid correlation Ophus refers to.
Works on images already in fourier space.
"""
m = img1_fft * conj(img2_fft)
M = sqrt(tf.abs(m))
magnitude = tf.complex(M, M * 0.0)
theta = angle(m)
euler = tf.exp(tf.complex(theta * 0.0, theta))
D = magnitude * euler
Icorr = real(ifft2(D))
return Icorr
def PST(I, LPF, Phase_strength, Warp_strength, Threshold_min, Threshold_max):
#inverting Threshold_min to simplyfy optimization porcess, so we can clip all variable between 0 and 1
LPF = ops.convert_to_tensor_v2(LPF)
Phase_strength = ops.convert_to_tensor_v2(Phase_strength)
Warp_strength = ops.convert_to_tensor_v2(Warp_strength)
I = ops.convert_to_tensor_v2(I)
Threshold_min = ops.convert_to_tensor_v2(Threshold_min)
Threshold_max = ops.convert_to_tensor_v2(Threshold_max)
Threshold_min = -Threshold_min
L = 0.5
x = tf.linspace(-L, L, I.shape[0])
y = tf.linspace(-L, L, I.shape[1])
[X1, Y1] = (tf.meshgrid(x, y))
X = tf.transpose(X1)
Y = tf.transpose(Y1)
[THETA, RHO] = cart2pol(X, Y)
# Apply localization kernel to the original image to reduce noise
Image_orig_f = sig.fft2d(tf.dtypes.cast(I, tf.complex64))
tmp6 = (LPF**2.0) / tfm.log(2.0)
tmp5 = tfm.sqrt(tmp6)
tmp4 = (tfm.divide(RHO, tmp5))
tmp3 = -tfm.pow(tmp4, 2)
tmp2 = tfm.exp(tmp3)
expo = fftshift(tmp2)
Image_orig_filtered = tfm.real(
sig.ifft2d((tfm.multiply(tf.dtypes.cast(Image_orig_f, tf.complex64),
tf.dtypes.cast(expo, tf.complex64)))))
# Constructing the PST Kernel
tp1 = tfm.multiply(RHO, Warp_strength)
PST_Kernel_1 = tfm.multiply(
tp1, tfm.atan(tfm.multiply(RHO, Warp_strength))
) - 0.5 * tfm.log(1.0 + tfm.pow(tf.multiply(RHO, Warp_strength), 2.0))
PST_Kernel = PST_Kernel_1 / tfm.reduce_max(PST_Kernel_1) * Phase_strength
# Apply the PST Kernel
temp = tfm.multiply(
fftshift(
tfm.exp(
tfm.multiply(tf.dtypes.complex(0.0, -1.0),
tf.dtypes.cast(PST_Kernel,
tf.dtypes.complex64)))),
sig.fft2d(tf.dtypes.cast(Image_orig_filtered, tf.dtypes.complex64)))
Image_orig_filtered_PST = sig.ifft2d(temp)
# Calculate phase of the transformed image
PHI_features = tfm.angle(Image_orig_filtered_PST)
out = PHI_features
out = (out / tfm.reduce_max(out)) * 3
return out
def train_bound(t):
"""Trains the model to equalize values and spatial derivatives at boundaries x=5
and x=-5 to enforce periodic boundary condition
Args:
t : A tf.Tensor of shape (batch_size,).
"""
x1 = 5 * tf.ones(t.shape)
x2 = -5 * tf.ones(t.shape)
with tf.GradientTape(True, False) as tape:
tape.watch(PINN.trainable_weights)
with tf.GradientTape(True, False) as grtape1:
grtape1.watch([t, x1, x2])
#Automatic differentiation of complex functions is weird in tensorflow
#so we differentiate real and imaginary parts seperately
h_real_1 = tfm.real(PINN(tf.stack([t, x1], -1)))
h_imag_1 = tfm.imag(PINN(tf.stack([t, x1], -1)))
h_real_2 = tfm.real(PINN(tf.stack([t, x2], -1)))
h_imag_2 = tfm.imag(PINN(tf.stack([t, x2], -1)))
#First order derivatives
h_x1_real = grtape1.gradient(h_real_1, x1)
h_x1_imag = grtape1.gradient(h_imag_1, x1)
h_x2_real = grtape1.gradient(h_real_2, x2)
h_x2_imag = grtape1.gradient(h_imag_2, x2)
#h1_real and h1_imag have shape (batch_size,2)
del grtape1
h1 = tf.complex(h_real_1, h_imag_1)
h1_x = tf.complex(h_x1_real, h_x1_imag)
h2 = tf.complex(h_real_2, h_imag_2)
h2_x = tf.complex(h_x2_real, h_x2_imag)
MSE = tfm.reduce_mean(
tfm.pow(tfm.abs(h1 - h2), 2) + tfm.pow(tfm.abs(h1_x - h2_x), 2))
grads = tape.gradient(MSE, PINN.trainable_weights)
sgd_opt.apply_gradients(zip(grads, PINN.trainable_weights))
return MSE
def train_colloc(t, x):
"""Trains the model to obey the given PDE at collocation points
Args:
t: A tf.Tensor of shape (batch_size,)
x: A tf.Tensor of shape (batch_size,).
"""
with tf.GradientTape(True, False) as tape:
tape.watch(PINN.trainable_weights)
#Calculate various derivatives of the output
with tf.GradientTape(True, False) as grtape0:
grtape0.watch([t, x])
with tf.GradientTape(True, False) as grtape1:
grtape1.watch([t, x])
#Automatic differentiation of complex functions is weird in tensorflow
#so we differentiate real and imaginary parts seperately
h_real = tfm.real(PINN(tf.stack([t, x], -1)))
h_imag = tfm.imag(PINN(tf.stack([t, x], -1)))
#First order derivatives
h_x_real = grtape1.gradient(h_real, x)
h_x_imag = grtape1.gradient(h_imag, x)
h_t_real = grtape1.gradient(h_real, t)
h_t_imag = grtape1.gradient(h_imag, t)
#h1_real and h1_imag have shape (batch_size,2)
del grtape1
#Second order derivatives
h_xx_real = grtape0.gradient(h_x_real, x)
h_xx_imag = grtape0.gradient(h_x_imag, x)
del grtape0
h = tf.complex(h_real, h_imag)
h_t = tf.complex(h_t_real, h_t_imag)
h_xx = tf.complex(h_xx_real, h_xx_imag)
j = tf.complex(0, 1)
MSE = tfm.reduce_euclidean_norm(
tfm.abs((j * h_t) + (0.5 * h_xx) + (tfm.conj(h) * h * h)))
grads = tape.gradient(MSE, PINN.trainable_weights)
sgd_opt.apply_gradients(zip(grads, PINN.trainable_weights))
del tape
return MSE
def cross_correlation_tf(A, B):
A = A - tf.math.reduce_mean(A)
B = B - tf.math.reduce_mean(B)
R = real(ifft2(fft2(A) * fft2(B[..., ::-1, ::-1])))
return R
def phase_correlation_tf(img1_fft, img2_fft):
"Perform phase correlation on images that are already in fourier space"
C = img1_fft * conj(img2_fft)
D = tf.abs(C)
return real(ifft2(C / tf.complex(D, D * 0.0)))