INNER CODE UNIT · Python
image_registration
chonyy/ML-auto-baseball-pitching-overlay · src/generate_overlay.py:73
def image_registration(ref_image, offset_image, shifts, list_idx, width, height):
# The shift is calculated once for each video and stored
if list_idx not in shifts:
xoff, yoff = cross_correlation_shifts(
ref_image[:, :, 0], offset_image.frame[:, :, 0]
)
shifts[list_idx] = (xoff, yoff)
else:
xoff, yoff = shifts[list_idx]
offset_image.ball = tuple(
[offset_image.ball[0] - int(xoff), offset_image.ball[1] - int(yoff)]
)
matrix = np.float32([[1, 0, -xoff], [0, 1, -yoff]])
corrected_image = cv2.warpAffine(offset_image.frame, matrix, (width, height))
return corrected_image