From e46e06d8401dbe5f5ae880ca1d9bd5073fa9c9ac Mon Sep 17 00:00:00 2001 From: Richard Bowman Date: Wed, 29 Apr 2020 12:00:13 +0100 Subject: [PATCH] add FFT tracking and "leapfrog" support --- .../camera_stage_tracker.py | 130 ++++++++++++-- .../fft_image_tracking.py | 168 ++++++++++++++++++ 2 files changed, 281 insertions(+), 17 deletions(-) create mode 100644 openflexure_microscope/api/default_extensions/camera_stage_mapping/fft_image_tracking.py diff --git a/openflexure_microscope/api/default_extensions/camera_stage_mapping/camera_stage_tracker.py b/openflexure_microscope/api/default_extensions/camera_stage_mapping/camera_stage_tracker.py index 9c7a7235..b643b69f 100644 --- a/openflexure_microscope/api/default_extensions/camera_stage_mapping/camera_stage_tracker.py +++ b/openflexure_microscope/api/default_extensions/camera_stage_mapping/camera_stage_tracker.py @@ -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) - \ No newline at end of file diff --git a/openflexure_microscope/api/default_extensions/camera_stage_mapping/fft_image_tracking.py b/openflexure_microscope/api/default_extensions/camera_stage_mapping/fft_image_tracking.py new file mode 100644 index 00000000..81b7efd0 --- /dev/null +++ b/openflexure_microscope/api/default_extensions/camera_stage_mapping/fft_image_tracking.py @@ -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)