add FFT tracking and "leapfrog" support

This commit is contained in:
Richard Bowman 2020-04-29 12:00:13 +01:00
parent ed8057ce04
commit e46e06d840
2 changed files with 281 additions and 17 deletions

View file

@ -11,6 +11,11 @@ import time
from numpy.linalg import norm
import cv2
from scipy import ndimage
from collections import namedtuple
import logging
from fft_image_tracking import high_pass_fft_template, displacement_from_fft_template, TrackingError
TrackerHistory = namedtuple("TrackerHistory", ["stage_positions", "image_positions"])
def central_half(image):
"""Return the central 50% (in X and Y) of an image"""
@ -25,7 +30,8 @@ def datum_pixel(image):
except:
return (np.array(image.shape[:2]) - 1) / 2.
def locate_feature_in_image(image, feature, margin=0, restrict=False):
########## Cross-correlation based tracking ############
def locate_feature_in_image(image, feature, margin=0, restrict=False, relative_to="top left"):
"""Find the given feature (small image) and return the position of its datum (or centre) in the image's pixels.
image : numpy.array
@ -39,6 +45,10 @@ def locate_feature_in_image(image, feature, margin=0, restrict=False):
restrict : bool (optional, default False)
If set to true, restrict the search area to a square of (margin * 2 + 1) pixels centred on the pixel that most
closely overlaps the datum points of the two images.
relative_to : string (optional, default "top left")
We return the position of the centre (or datum pixel, if it's got that metadata) of the feature, relative to
either the top left (i.e. 0,0) pixel in the image, or the central pixel - to do the latter, set ``relative_to``
to "centre" (or "center" if you must).
The `image` must be larger than `feature` by a margin big enough to produce a meaningful search area. We use the
OpenCV `matchTemplate` method to find the feature. The returned position is the position, relative to the corner of
@ -69,21 +79,55 @@ def locate_feature_in_image(image, feature, margin=0, restrict=False):
assert np.sum(corr) > 0, "Error: the correlation image doesn't have any nonzero pixels."
peak = ndimage.measurements.center_of_mass(corr) # take the centroid (NB this is of grayscale values, not binary)
pos = np.array(peak) + image_shift + datum_pixel(feature) # return the position of the feature's datum point.
return pos
if relative_to in ["top left", None]:
return pos
if relative_to in ["centre", "center"]:
return pos - (np.array(image.shape[:2]) - 1)/2.
raise ValueError("An invalid value was specified for datum.")
######## FFT based tracking functions ###########
# moved to fft_image_tracking.py
class Tracker():
def __init__(self, grab_image, get_position, settle=None):
def __init__(self, grab_image, get_position, settle=None, method="direct", **kwargs):
"""A class to manage moving the stage and following motion in the image
Constructor Arguments:
grab_image: a function that returns the image as a numpy array
get_position: a function that returns position as a numpy array
settle: a function that waits and/or discards images
method: string, currently "direct" (default) or "fft"
Additional keyword arguments are passed to the tracking method.
Tracking methods:
"direct": uses cross-correlation with the central half of the image.
There are currently no options for this method.
"fft" uses FFT-based cross-correlation, after a high pass filter.
Keyword arguments are accepted:
pad: boolean, default=True
Whether to zero-pad the FFT to remove ambiguity. If
false, the answer is only unique modulo one field of
view due to the periodic nature of FFTs. Setting this
option to False speeds up tracking by about 4x.
sigma: floating point, default=10
The standard deviation, in pixels, of a Gaussian filter
used to smooth the image, before subtracting the smooth
image from the original, in a low pass filter. NB the
standard deviation is given in pixels, but is applied
in the Fourier domain (with appropriate transformation).
The value of sigma does not affect computation speed.
We accept functions because that seems like the easiest way to be
compatible with many different cameras/stages. Subclass and override
``__init__`` if you want to use a particular object instead.
# Subclassing notes
If you change the tracking method, you should override:
* track_image
* generate_template
* max_displacement
* min_displacement (optional - defaults to -max_displacment)
NB the ``image_position`` that this class returns may be the negative of
what you might expect. This is because normally we are looking for
where a certain object (usually matched to a template image) is within
@ -97,7 +141,11 @@ class Tracker():
self._template = None
self.margin = np.array([0, 0])
self._template_position = np.array([0.0, 0.0])
self._last_point = None
self.image_shape = None
self.method = method
#self.kwargs = {"error_threshold": 0.2}.update(kwargs)
self.kwargs = kwargs
def get_position(self):
"""Get the position of the stage"""
@ -138,42 +186,86 @@ class Tracker():
Whether to wait for the stage to settle before taking the template image
reset_history: bool, default True
Whether to erase all the previously-stored positions
relative_positions: bool, default True
If true, we will define the first point (as read from the camera) to be [0,0]
and make all future measurements relative to this one. NB this won't affect
the stage positions, which are always absolute.
"""
if settle:
self.settle()
image = self._grab_image()
self.template = central_half(image)
self.template = self.generate_template(image)
self.image_shape = image.shape
self.margin = np.array(image.shape)[:2] - np.array(self.template.shape)[:2]
if reset_history:
self.reset_history()
self._template_position = np.array([0., 0.])
if relative_positions:
self._template_position = self.track_image(image) # Position should be zero initially
self.append_point(settle=False)
def leapfrog(self):
"""Replace the template but don't change position.
By default, this will replace the template with the image from the last
point we measured, but update _template_position so that the coordinates
returned don't change.
"""
if self._last_point is None:
raise ValueError("Can't leapfrog until you have measured at least one point.")
image, image_pos = self._last_point
self.template = self.generate_template(image)
if np.any(self.image_shape != image.shape):
raise ValueError("Error: the image size seems to have changed!")
# Ensure that the position doesn't change. NB this is nice and reliable because
# it actually runs self.track_image, but we could be much more efficient if we
# assumed that self.track_image(image) == 0, and we can certainly do the maths
# to make that work...
# TODO: eliminate the unnecessary correlation
self._template_position = np.array(image_pos)
def generate_template(self, image):
"""Generate a template based on a supplied image.
This function is designed to be overridden in order to
change the tracking method.
"""
if self.method == "direct":
return central_half(image)
if self.method == "fft":
kwargs = {k: v for k, v in self.kwargs.items() if k in ["pad", "sigma"]}
return high_pass_fft_template(image, calculate_peak=True, **kwargs)
@property
def max_displacement(self):
"""The highest position values that can be tracked"""
return self.margin // 2 # TODO: be cleverer about non-trivial values of template_position
if self.method == "direct":
# TODO: if template_position is not central, should we alter this??
disp = (np.array(self.image_shape[:2]) - np.array(self.template.shape)[:2]) // 2
if self.method == "fft":
# FFT tracking does a real FFT to track the position, which is half as long in
# the last dimension. If we didn't zero pad, the transform will have the same shape as
# the image in x, and half in y - so we return half the image size. If we are zero
# padding, then both these dimensions double, and we return the image size.
disp = np.array(self.template.shape) // np.array([2,1])
return disp + self._template_position
@property
def min_displacement(self):
"""The lowest position values that can be tracked"""
return -self.max_displacement # TODO: be cleverer about non-trivial template_position values
return self._template_position - self.max_displacement
# TODO: be cleverer about tracking assymetry? Currently there is none...
@property
def max_safe_displacement(self):
"""The biggest displacement we can safely attempt to track without knowing direction."""
return np.min(np.concatenate([self.max_displacement, -self.min_displacement]))
return np.min(np.concatenate([self.max_displacement - self._template_position,
-self.min_displacement - self._template_position]))
def point_in_safe_range(self, point):
"""Return True if a given point is within the safe range of the tracker."""
return np.all(point > self.min_displacement) and np.all(point < self.max_displacement)
def track_image(self, image):
"""Find the position of the image relative to the template
This uses the method specified at initialisation time to
track motion of the sample.
NB this class is intended to track motion of the sample - most of
the time, we're interested in the motion of a (small) object that
is represented by the template, relative to the (larger) image. In
@ -182,7 +274,11 @@ class Tracker():
a minus sign in front of `locate_feature_in_image` in the source
code.
"""
return - locate_feature_in_image(image, self.template) - self._template_position
if self.method=="direct":
return -locate_feature_in_image(image, self.template, relative_to="centre") + self._template_position
if self.method=="fft":
kwargs = {k: v for k, v in self.kwargs.items() if k in ["pad", "fractional_threshold", "error_threshold"]}
return -displacement_from_fft_template(self.template, image, **kwargs) + self._template_position
def append_point(self, settle=True, image=None):
"""Find the current position using both stage and image, and append it"""
@ -194,6 +290,7 @@ class Tracker():
stage_pos = self.get_position()
self._image_positions.append(image_pos)
self._stage_positions.append(stage_pos)
self._last_point = (image, image_pos)
return stage_pos, image_pos
@property
@ -209,7 +306,7 @@ class Tracker():
@property
def history(self):
"""Return arrays of stage, image positions"""
return self.stage_positions, self.image_positions
return TrackerHistory(self.stage_positions, self.image_positions)
def reset_history(self, leave_first_point=False):
"""Reset the positions and displacements recorded"""
@ -286,4 +383,3 @@ def concatenate_tracker_histories(histories):
"""
components = zip(*histories)
return tuple(np.concatenate(c, axis=1) for c in components)

View file

@ -0,0 +1,168 @@
"""
Utility functions to track motion of a microscope using FFT-based correlation.
Cross-correlation is a reasonable way to determine where an object is in an
image. It can also be used to track 2D motion. The Fourier Shift Theorem
relies on the fact that a correlation (or convolution) becomes a multiplication
in the Fourier domain. This means that Fast Fourier Transforms are an
efficient way to implement cross-correlation of whole images. This module
contains a number of functions to simplify tracking the motion of a microscope
stage using FFTs.
(c) Richard Bowman 2020, released under GNU GPL v3
No warranty, express or implied, is given with respect to this code.
"""
import numpy as np
import logging
from array_with_attrs import ArrayWithAttrs
def grayscale_and_padding(image, pad=True):
"""Convert to grayscale and prepare for zero padding if needed.
The FFT-based tracking methods need grayscale images. Also, if
we are going to zero-pad, we should convert to floating point and
ensure the mean of the image is zero, otherwise the dominant feature
will be the edge of the image.
Returns:
image, fft_shape
"""
if len(image.shape) == 3:
image = np.mean(image, axis=2)
fft_shape = np.array(image.shape)
if pad:
image = image.astype(np.float) - np.mean(image)
fft_shape *= 2
return image, fft_shape
def high_pass_fourier_mask(shape, s, rfft=True):
"""Generate a mask performing a high pass filter
The return value is a 2D array, which can be multiplied
with the Fourier Transform of an image to perform a high
pass filter.
Arguments:
shape: tuple of 2 integers
The shape of the output array
s: float
The standard deviation of the Gaussian in real
space, in pixels
"""
high_pass_filter = np.ones(shape)
x, y = (np.arange(n, dtype=np.float) for n in shape)
# Beyond the halfway point of the array, frequencies are negative
x[x.shape[0]//2:x.shape[0]] -= x.shape[0]
if not rfft: # If it's a real fft, the last axis is halved so we can skip this.
y[y.shape[0]//2:y.shape[0]] -= y.shape[0]
x /= np.max(np.abs(x)) * 2 # Normalise so highest frequency is 1/2
y /= np.max(np.abs(y)) * 2 # Normalise so highest frequency is 1/2
r2 = x[:, np.newaxis]**2 + y[np.newaxis, :]**2
# now we multiply by 1-FT(Gaussian kernel with sd of s pixels)
high_pass_filter -= np.exp(-2*np.pi**2*s**2*r2)
return high_pass_filter
def high_pass_fft_template(image, sigma=10, pad=True, calculate_peak=True):
"""Calculate a high-pass-filtered FT template for tracking
This performs a real FFT, and then attenuates low frequencies.
The resulting array can be used as a template for tracking.
sigma is the standard deviation in pixels of the Gaussian used
in the high pass filter.
pad enables (default) zero padding - this removes the ambiguity
around position, at the cost of making the function slower. We
subtract the mean and zero-pad the input array (equivalent to
padding with the mean value, to reduce the impact of the edge)
calculate computes the value of the brightest pixel we'd expect
in a correlation image (i.e. the peak if we correlate the image
passed in with the template we're generating). This is stored
in ``template.attrs["peak_correlation_value"]``
"""
image, fft_shape = grayscale_and_padding(image, pad)
initial_fft = np.fft.rfft2(image, s=fft_shape) # NB rfft2 is faster, but a different shape!
high_pass_filter = high_pass_fourier_mask(initial_fft.shape, sigma)
if calculate_peak:
expected_peak = np.mean(np.conj(initial_fft) * high_pass_filter * initial_fft)
template = ArrayWithAttrs(np.conj(initial_fft) * high_pass_filter)
template.attrs["maximum_correlation_value"] = expected_peak
return template
return np.conj(initial_fft) * high_pass_filter
def background_subtracted_centre_of_mass(corr, fractional_threshold=0.05, quadrant_swap=False):
"""Carry out a background subtracted centre of mass measurement
Arguments:
corr: a 2D numpy array, to be thresholded
fractional_threshold: the fraction of the range (from
min(corr) to max(corr)) that should remain above
the background level. 1 means no thresholding,
0.05 means use only the top 5% of the range.
quadrant_swap: boolean, default False
Set this to true if we are working on the output of
a Fourier transform. This will adjust the coordinates
such that we effectively perform quadrant swapping, to
place the DC component in the centre of the image, and
make the coordinate (0,0) correspond to that point, with
positive and negative coordinates either side.
"""
assert corr.dtype == np.float, "The image must be floating point"
background = np.max(corr) - fractional_threshold * (np.max(corr) - np.min(corr))
background_subtracted = corr - background
background_subtracted[background_subtracted < 0] = 0
xs, ys = (np.arange(n) for n in corr.shape) # This is equivalent to meshgrid, more or less...
if quadrant_swap:
xs[len(xs)//2:] -= len(xs)
ys[len(ys)//2:] -= len(ys)
x = np.sum(background_subtracted * xs[:, np.newaxis])
y = np.sum(background_subtracted * ys[np.newaxis, :])
I = np.sum(background_subtracted)
return np.array([x/I, y/I])
class TrackingError(Exception):
pass
def displacement_from_fft_template(template, image, fractional_threshold=0.1, pad=True, return_peak=False, error_threshold=0):
"""Find the displacement, in pixels, of an image from a template
The template should be generated by ``high_pass_fft_template``
Fractional_threshold is the fraction of the range (from max to min)
of the cross-correlation image that should remain above the threshold
before finding the peak by centre-of-mass.
NB because of the periodic boundary conditions of the FFT, this gives
a result that is ambiguous - it's only accurate modulo one image.
The result that is returned represents the smalles displacement,
positive or negative. You may add or subtract one whole image-width
(or height) if that makes sense - use other cues to resolve the
ambiguity.
return_peak returns the brightes pixel in the correlation image, as
well as the displacement in a tuple.
error_threshold is an optional floating-point number between 0 and 1.
Setting it to a value greater than 0 will compare the correlation value
with the maximum possible. If the ratio of the current signal to the
maximum drops below ``error_threshold``, we raise a ``TrackingError``
exception.
"""
image, fft_shape = grayscale_and_padding(image, pad)
# The template is already Fourier transformed and high pass filtered.
# so multiplying the two in Fourier space performs the convolution.
corr = np.fft.irfft2(template * np.fft.rfft2(image, s=fft_shape))
if error_threshold > 0:
if np.max(corr)/template.attrs["maximum_correlation_value"] < error_threshold:
raise TrackingError("The correlation signal dropped below the threshold set.")
displacement = background_subtracted_centre_of_mass(corr, fractional_threshold, quadrant_swap=True)
if return_peak:
return displacement, np.max(corr)
return displacement
def displacement_between_images(image_0, image_1, sigma=10, fractional_threshold=0.1, pad=True):
"""Calculate the displacement, in pixels, between two images."""
return displacement_from_fft_template(high_pass_fft_template(image_0, sigma, pad=pad),
image_1, fractional_threshold, pad=pad)