import numpy as np
import scipy
import pyklip.klip
def _img_xcorr_shift(shift, frame1, frame2, mask=None):
"""
Compute the negative of the xcorr between frame1 and frame2 given shift:
Args:
shift: [dx, dy] in pixels to shift frame2
frame1: refernce frame of shape (Y, X)
frame2: frame to align of shape (Y, X)
mask: shape of (Y, X) where =1 for good pixels and =0 for bad pixels
Return:
negative of xcorr between frame1 and the shifted frame2
"""
dx, dy = shift
oldcenter = [frame2.shape[1], frame2.shape[0]]
frame2_shifted = pyklip.klip.align_and_scale(frame2, [oldcenter[0] + dx, oldcenter[1] + dy], oldcenter)
if mask is not None:
good = np.where(mask > 0)
frame1 = frame1[good]
frame2_shifted = frame2_shifted[good]
return -np.nansum(frame1*frame2_shifted)
[docs]
def find_best_shift(frame, ref_frame, guess_center=None, inner_mask=5, outer_mask=None):
"""
Finds the best shift based on xcorr
Args:
frame: frame to find offset of. Shape of (Y, X)
ref_frame: reference frame to align frame to. Shape of (Y, X)
Return:
(dx, dy): best shift to shift frame to match reference frame
"""
if guess_center is None:
guess_center = np.array([frame.shape[1]/2., frame.shape[0]/2.])
if outer_mask is None:
outer_mask = frame.shape[0]/2 - 10
# make mask
y, x = np.indices(frame.shape)
r = np.sqrt((x-guess_center[0])**2 + (y-guess_center[1])**2)
mask = np.ones(frame.shape)
bad = np.where((r < inner_mask) | (r > outer_mask))
mask[bad] = 0
result = scipy.optimize.minimize(_img_xcorr_shift, (0,0), args=(ref_frame, frame, mask), method="Nelder-Mead")
return result.x