diff --git a/test/test_registration.py b/test/test_registration.py index 61f92819..2acecc29 100644 --- a/test/test_registration.py +++ b/test/test_registration.py @@ -2,11 +2,12 @@ import tempfile from test_utils import PySparkTestCase -from numpy import allclose, dstack, mean, random, ndarray +from numpy import allclose, dstack, mean, random, ndarray, zeros, pi from scipy.ndimage.interpolation import shift from nose.tools import assert_true from thunder.registration.registration import Registration +from thunder.registration.transformation import TranslationTransformation, EuclideanTransformation from thunder.rdds.fileio.imagesloader import ImagesLoader @@ -173,3 +174,51 @@ def test_crosscorrVolume(self): imOut = Registration('planarcrosscorr').prepare(ref).run(imIn).first()[1] assert_true(allclose(paramOut, [[2, -2], [2, -2], [2, -2]])) assert_true(allclose(ref[:-2, 2:, :], imOut[:-2, 2:, :])) + + +class TestTransformation(ImgprocessingTestCase): + def test_translation(self): + im = zeros((5, 5)) + im[2:4, 2:4] = 1.0 + expected = zeros((5, 5)) + expected[3:, :2] = 1.0 + tfm = TranslationTransformation(shift=(1, -2)) + out = tfm.apply(im) + assert_true(allclose(out, expected)) + + def _run_euclidean_test(self, sz): + im = zeros((sz, sz)) + im[:1, :] = 1.0 + expected = zeros((sz, sz)) + expected[:, 2] = 1.0 + tfm = EuclideanTransformation(shift=(2.0, 0.0), rotation=pi/2) + out = tfm.apply(im) + assert_true(allclose(out, expected)) + + def test_euclidean(self): + self._run_euclidean_test(5) + self._run_euclidean_test(6) + + +class TestLucasKanade(ImgprocessingTestCase): + def test_euclidean(self): + vol = zeros((16, 16)) + vol[4:12, 4:8] = 1.0 + rot = pi / 8.0 + trueTfm = EuclideanTransformation((0, 1), rot) + volTfm = trueTfm.apply(vol) + + reg = Registration('lucaskanade', transformationType='Euclidean').prepare(volTfm) + predTfm = reg.getTransform(vol) + print trueTfm.getParams(), predTfm.getParams() + assert_true(allclose(trueTfm.getParams(), predTfm.getParams(), atol=1e-2)) + + def test_translation(self): + vol = zeros((16, 16)) + vol[4:12, 4:8] = 1.0 + trueTfm = TranslationTransformation((1.75, -2.5)) + volTfm = trueTfm.apply(vol) + reg = Registration('lucaskanade', transformationType='Translation').prepare(volTfm) + predTfm = reg.getTransform(vol) + print trueTfm.getParams(), predTfm.getParams() + assert_true(allclose(trueTfm.getParams(), predTfm.getParams(), atol=1e-2)) diff --git a/thunder/registration/__init__.py b/thunder/registration/__init__.py index f5c503ec..7248d229 100644 --- a/thunder/registration/__init__.py +++ b/thunder/registration/__init__.py @@ -1 +1,2 @@ -from thunder.registration.methods.crosscorr import CrossCorr, PlanarCrossCorr \ No newline at end of file +from thunder.registration.methods.crosscorr import CrossCorr, PlanarCrossCorr +from thunder.registration.methods.lucaskanade import LucasKanade diff --git a/thunder/registration/methods/lucaskanade.py b/thunder/registration/methods/lucaskanade.py new file mode 100644 index 00000000..ba91ea55 --- /dev/null +++ b/thunder/registration/methods/lucaskanade.py @@ -0,0 +1,94 @@ +""" Registration methods based on Lucas-Kanade registration""" + +from numpy import array, ndarray, inf, zeros + +from thunder.rdds.images import Images +from thunder.registration.registration import RegistrationMethod +from thunder.registration.transformation import TRANSFORMATION_TYPES +from thunder.registration.methods.utils import volumesToMatrix, solveLinearized, computeReferenceMean, checkReference + + +class LucasKanade(RegistrationMethod): + """Lucas-Kanade registration method. + + Lucas-Kanade (LK) is an iterative algorithm for aligning an image to a reference. It aims to minimize the squared + error between the transformed image and the reference. As the relationship between transformation parameters and + transformed pixels is nonlinear, we have to perform a series of linearizations similar to Levenberg-Marquardt. + At each iteration, we compute the Jacobian of the output image with respect to the input parameters, and solve a + least squares problem to identify a change in parameters. We update the parameters and repeat. + + To increase robustness, we extend the traditional LK algorithm to use a set of reference images. + We minimize the squared error of the difference of the transformed image and a learned weighting of the references. + """ + + def __init__(self, transformationType='Translation', border=0, tol=1e-5, maxIter=100, robust=False): + """ + Parameters + ---------- + transformationType : one of 'Translation', 'Euclidean', optional, default = 'Translation' + type of transformation to use + border : int or tuple, optional, default = 0 + Border to be zeroed out after transformations. For most datasets, it is + critical that this value be larger than the maximum translation to get + good results and avoid boundary artifacts. + maxIter : int, optional, default = 100 + maximum number of iterations + tol : float, optional, default = 1e-5 + stopping criterion on the L2 norm of the change in parameters + robust : bool, optional, default = False + solve a least absolute deviation problem instead of least squares + """ + self.transformationType = transformationType + self.border = border + self.maxIter = maxIter + self.tol = tol + self.robust = robust + + def prepare(self, images, startidx=None, stopidx=None): + """ + Prepare Lucas-Kanade registration by computing or specifying a reference image. + + Parameters + ---------- + images : ndarray or Images object + Images to compute reference from, or a single image to set as reference + + See computeReferenceMean. + """ + if isinstance(images, Images): + self.reference = computeReferenceMean(images, startidx, stopidx) + elif isinstance(images, ndarray): + self.reference = images + else: + raise Exception('Must provide either an Images object or a reference') + # Convert references to matrix to speed up solving linearized system + self.referenceMat = volumesToMatrix(self.reference) + return self + + def isPrepared(self, images): + """ + Check if Lucas-Kanade is prepared by checking the dimensions of the reference. + + See checkReference. + """ + + if not hasattr(self, 'reference'): + raise Exception('Reference not defined') + else: + checkReference(self.reference, images) + + def getTransform(self, vol): + from numpy.linalg import norm + # Create initial transformation + tfm = TRANSFORMATION_TYPES[self.transformationType](shift=zeros(vol.ndim)) + iter = 0 + normDelta = inf + params = [] + while iter < self.maxIter and normDelta > self.tol: + volTfm, jacobian = tfm.jacobian(vol, border=self.border) + deltaTransformParams, coeff = solveLinearized(volumesToMatrix(volTfm), volumesToMatrix(jacobian), self.referenceMat, self.robust) + tfm.updateParams(deltaTransformParams) + normDelta = norm(deltaTransformParams) + params.append(tfm.getParams()) + iter += 1 + return tfm diff --git a/thunder/registration/methods/utils.py b/thunder/registration/methods/utils.py index 26571f95..40c3127d 100644 --- a/thunder/registration/methods/utils.py +++ b/thunder/registration/methods/utils.py @@ -122,4 +122,141 @@ def computeDisplacement(arry1, arry2): # cast to basic python int for serialization adjusted = [int(d - n) if d > n // 2 else int(d) for (d, n) in pairs] - return adjusted \ No newline at end of file + return adjusted + +def zeroBorder(vol, border=1, cval=0.0): + """Zero-out boundary voxels of a volume. + + Parameters + ---------- + vol: 3d volume + border: scalar or tuple containing borders for each dimension + cval: constant value to replace boundary voxels with, default is 0.0 + """ + import numpy as np # sorry + vol = vol.copy() # create new copy so we don't overwrite vol + dims = np.array(vol.shape) + if np.size(border) == 1: + border = border * np.ones(vol.ndim, dtype=int) + border = np.array(border, dtype=int) + # Don't apply border to singleton dimensions. + border[dims == 1] = 0 + assert len(border) == vol.ndim + if np.any(dims - border <= 0): + raise ValueError('Border %s exceeds volume shape %s.' % + (str(border), str(dims)) ) + for dim, bval in enumerate(border): + if bval > 0: + slices = [slice(bval) if d == dim else slice(None) for d in xrange(vol.ndim)] + vol[slices] = cval + slices[dim] = slice(-bval, None) + vol[slices] = cval + return vol + + +def solveLinearized(vec, jacobian, reference, robust=False): + """Solve linearized registration problem for change in transformation parameters and weighting on reference basis. + + Parameters + ---------- + vec : array, shape (nvoxels,) + vectorized volume + jacobian: array, shape (nvoxels, nparams) + jacobian for each transformation parameter + reference: array, shape (nvoxels, nbasis) + array of vectorized reference volumes + robust: bool, optional, default = False + solve a least absolute deviation problem instead of least squares + + Returns + ------- + deltaTransforms : array, shape (nparams,) + optimal change in transformation parameters + coeff : array, shape (nbasis,) + optimal weighting of reference volumes + """ + from numpy import column_stack + A = column_stack((jacobian, reference)) + if robust: + from statsmodels.regression.quantile_regression import QuantReg + quantile = 0.5 + model = QuantReg(vec, A).fit(q=quantile) + params = model.params + else: + from numpy.linalg import lstsq + model = lstsq(A, vec) + params = model[0] + from numpy import split + deltaTransform, coeff = split(params, [jacobian.shape[1]]) + return deltaTransform, coeff + + +def volToVec(vol): + """ + Convert volume to vector. + + Parameters + ---------- + vol : array + volume + + Returns + ------- + vectorized volume + """ + return vol.ravel() + + +def vecToVol(vec, dims): + """ + Convert vector to volume. + + Parameters + ---------- + vec : array, shape (nvoxels,) + vectorized volume + dims : tuple, optional + shape of volume, must be set to reshape vectors into volumes + + Returns + ------- + volume if input is a vector, and vector if input is a volume + """ + assert vec.size == prod(dims) + return vec.reshape(dims) + + +def volumesToMatrix(vols): + """Convert list of volumes to a matrix. + + Parameters + ---------- + vols : list of arrays + + Returns + ------- + array with size nvoxels by number of volumes + """ + from numpy import column_stack + if not isinstance(vols, list): + return volToVec(vols) + else: + return column_stack(map(volToVec, vols)).squeeze() + + +def imageGradients(im, sigma=None): + """Compute gradients of volume in each dimension using a Sobel filter. + + Parameters + ---------- + im : ndarray + single volume + sigma : float or tuple, optional, default = None + smoothing amount to apply to volume before computing gradients + """ + from scipy.ndimage.filters import gaussian_filter, sobel + if sigma is not None: + im = gaussian_filter(im, sigma) + grads = [sobel(im, axis=dim, mode='constant') / 8.0 for dim in xrange(im.ndim)] + return grads + diff --git a/thunder/registration/registration.py b/thunder/registration/registration.py index 9460555c..5fa2f20f 100644 --- a/thunder/registration/registration.py +++ b/thunder/registration/registration.py @@ -20,15 +20,17 @@ class Registration(object): def __new__(cls, method, **kwargs): from thunder.registration.methods.crosscorr import CrossCorr, PlanarCrossCorr + from thunder.registration.methods.lucaskanade import LucasKanade REGMETHODS = { 'crosscorr': CrossCorr, - 'planarcrosscorr': PlanarCrossCorr + 'planarcrosscorr': PlanarCrossCorr, + 'lucaskanade': LucasKanade, } checkParams(method, REGMETHODS.keys()) - return REGMETHODS[method](kwargs) + return REGMETHODS[method](**kwargs) @staticmethod def load(file): diff --git a/thunder/registration/transformation.py b/thunder/registration/transformation.py index 849375e4..c6500abb 100644 --- a/thunder/registration/transformation.py +++ b/thunder/registration/transformation.py @@ -1,4 +1,5 @@ """ Transformations produced by registration methods """ +import numpy as np from numpy import asarray from thunder.utils.serializable import Serializable @@ -9,6 +10,297 @@ class Transformation(object): def apply(self, im): raise NotImplementedError +# Module-level dictionary cache mapping volume dimensions to an array containing +# the image coordinates for that volume. These coordinates are used to speed up +# application of transformations that inherit from GridMixin. +_grid = {} + +def getGrid(dims): + """ + Check and (if needed) initialize a grid with the appropriate dimensions. + + Parameters + ---------- + dims : tuple + shape of volume + """ + global _grid + if dims not in _grid: + _grid[dims] = np.vstack((np.array(np.meshgrid(*[np.arange(d) for d in dims], indexing='ij')), + np.ones(dims)[np.newaxis, :, :])) + # Prevent grid from being altered. + _grid[dims].flags.writeable = False + + return _grid[dims] + +class GridMixin(object): + def getCoords(self, grid): + """ + Get the coordinates where the input volume is evaluated. + + Returns + ------- + array with size ndims by nvoxels specifying the coordinates for each voxel. + """ + raise NotImplementedError + + def apply(self, vol, order=1, cval=0.0, **kwargs): + from scipy.ndimage import map_coordinates + grid = getGrid(vol.shape) + coords = self.getCoords(grid) + tvol = map_coordinates(vol, coords, order=order, cval=0.0, **kwargs) + return tvol + +class DifferentiableTransformation(Transformation): + """ + Differentiable transformations must have methods to compute the Jacobian and update parameters. These are used by + iterative alignment techniques like Lucas-Kanade and RASL. + """ + + def _jacobianCoords(self, vol, sigma=None, normalize=True, border=1): + """Compute Jacobian of volume with respect to the coordinates. + + Args: + vol: volume + tfm: Transform object + sigma: smoothing bandwidth for gradients (None for no smoothing) + normalize: Whether to normalize images before aligning. + border: Number or tuple of border sizes to zero after transforming. + Returns: + tvol : array + transformed volume + grads : list of arrays + list of volume Jacobians with respect to coordinates, one for each dimension + """ + from numpy.linalg import norm + from thunder.registration.methods.utils import imageGradients, zeroBorder + + grads = imageGradients(vol, sigma) + tvol = zeroBorder(self.apply(vol), border) + normVol = norm(tvol.ravel()) + if normVol == 0.0: + raise ValueError('Transform yields volume of zeroes.') + grads = [zeroBorder(self.apply(grad), border) for grad in grads] + if normalize: + if normVol == 0.0: + raise ValueError('Transform yields volume of zeroes.') + # Update gradients to reflect normalization + grads = [grad / normVol - (grad * tvol).sum() / (normVol**3) * tvol for grad in grads] + tvol /= normVol + return tvol, grads + + def jacobian(self, vol, **kwargs): + """ + Compute gradient of transformation output with respect to parameters. + + Parameters + ---------- + vol : array + volume + + Returns + ------- + ordered list of arrays containing the Jacobian for each parameter + """ + raise NotImplementedError + + def getParams(self): + """ + Get the parameters of this transformation. + + Returns + ------- + array with shape (nparams,) + """ + raise NotImplementedError + + def setParams(self, params): + """ + Set each parameter in the transformation. + + Parameters + ---------- + params : array, shape (nparams,) + new values of parameters. The ordering here must match + the ordering of the parameters returned by jacobian. + """ + raise NotImplementedError + + def updateParams(self, deltaParams): + """ + Update each parameter in the transformation. + + Parameters + ---------- + deltaParams : array, shape (nparams,) + ordered changes in parameters to be applied. The ordering here must match + the ordering of the parameters returned by jacobian. + """ + self.setParams(self.getParams() + deltaParams) + + +class TranslationTransformation(DifferentiableTransformation): + def __init__(self, shift): + """Translation in 3d. + + Parameters + ---------- + shift: 3d vector of translations (x, y, z) + + """ + self.shift = shift + + def jacobian(self, vol, **kwargs): + return self._jacobianCoords(vol, **kwargs) + + def getParams(self): + return np.asarray(self.shift) + + def setParams(self, shift): + assert len(shift) == len(self.shift) + self.shift = shift + + def apply(self, vol): + from scipy.ndimage.interpolation import shift + return shift(vol, self.shift, mode='constant', cval=0.0) + + +class ProjectiveTransformation(GridMixin, DifferentiableTransformation): + """ + Projective transformations are differentiable and can be represented as a matrix.""" + + def __init__(self, center=None): + self.center = center + + def asMatrix(self): + raise NotImplementedError + + def getCoords(self, grid): + d = grid.shape[0] - 1 # number of dimensions + dims = grid.shape[1:] + if self.center is None: + self.center = (np.array(dims) - 1) / 2.0 + + A = self.asMatrix() + + # Center grid so that we apply rotations w.r.t. given center + center_bcast = np.r_[self.center, 0].reshape((-1, ) + (1,) * d) + if self.center is not None: + grid = grid - center_bcast + + # Apply transformation + grid = np.tensordot(A.T, grid, axes=(0,0)) + + # Move back to voxel reference frame + if self.center is not None: + grid = grid + center_bcast + return grid[:-1] # throw out last homogeneous coordinate + +class EuclideanTransformation(ProjectiveTransformation): + def __init__(self, shift, rotation=None, zTranslation=False, zRotation=False, center=None): + """Translation and rotation in 3d. + + Parameters + ---------- + shift : list or array + spatial shifts for each dimension + rotation : float, list, or array + rotation in x-y if scalar. rotation in x-y plane, x-z plane, y-z plane if list + or array and dataset is in 3d. + zTranslation : bool, optional, default = False + whether to allow translation in z + zRotation : bool, optional, default = False + whether to allow rotation in z + center : list or array, optional, default = None + Set center coordinates that define the point around which the volume is rotated. + Coordinates should be in terms of zero-indexed pixels. For example if the volume was 5x5x3, + the default center of rotation would be at the center of the volume: (2, 2, 1). + """ + self.shift = np.atleast_1d(shift) + self.ndim = len(self.shift) + if rotation is None: + if self.ndim == 2 or not zRotation: + rotation = 0.0 + else: + rotation = np.zeros(3) + self.rotation = np.atleast_1d(rotation) + self.zRotation = zRotation + self.zTranslation = zTranslation + self.center = center + + def getParams(self): + return np.r_[self.shift, self.rotation] + + def updateParams(self, deltaParams): + if self.ndim == 2: + self.shift += deltaParams[:2] + self.rotation += deltaParams[2] + else: + self.shift += deltaParams[:3] + self.rotation += deltaParams[3:] + + def jacobian(self, vol, **kwargs): + + tvol, imageGradients = self._jacobianCoords(vol, **kwargs) + imageGrid = getGrid(vol.shape) + ndim = len(self.shift) + + stheta = np.sin(self.rotation[0]) + ctheta = np.cos(self.rotation[0]) + + if ndim == 2 or not self.zRotation: + dtheta = imageGradients[0] * (-imageGrid[0] * stheta + imageGrid[1] * -ctheta) + dtheta += imageGradients[1] * ( imageGrid[0] * ctheta + imageGrid[1] * -stheta) + dangles = [dtheta] + else: + sphi = np.sin(self.rotation[1]) + cphi = np.cos(self.rotation[1]) + spsi = np.sin(self.rotation[2]) + cpsi = np.cos(self.rotation[2]) + dtheta = imageGradients[0] * ( + -imageGrid[0] * cphi * stheta + + imageGrid[1] * (-ctheta * cpsi - stheta * sphi * spsi) + + imageGrid[2] * (-cpsi * stheta * sphi + ctheta * spsi)) + dtheta += imageGradients[1] * ( + imageGrid[0] * ctheta * cphi + + imageGrid[1] * (-cpsi * stheta + ctheta * sphi * spsi) + + imageGrid[2] * (ctheta * cpsi * sphi + stheta * spsi)) + dphi = imageGradients[0] * ( + imageGrid[2] * ctheta * cphi * cpsi - + imageGrid[0] * ctheta * sphi + + imageGrid[1] * ctheta * cphi * spsi) + dphi += imageGradients[1] * ( + imageGrid[2] * cphi * cpsi * stheta - + imageGrid[0] * stheta * sphi + + imageGrid[1] * cphi * stheta * spsi) + dphi += imageGradients[2] * ( + -imageGrid[0] * cphi - + imageGrid[2] * cpsi * sphi - + imageGrid[1] * sphi * spsi) + dpsi = imageGradients[0] * ( + imageGrid[1] * (ctheta * cpsi * sphi + stheta * spsi) + + imageGrid[2] * (cpsi * stheta - ctheta * sphi * spsi) ) + dpsi += imageGradients[1] * ( + imageGrid[1] * (cpsi * stheta * sphi - ctheta * spsi) + + imageGrid[2] * (-ctheta * cpsi - stheta * sphi * spsi)) + dpsi += imageGradients[2] * (imageGrid[1] * cphi * cpsi - imageGrid[2] * cphi * spsi) + dangles = [dtheta, dphi, dpsi] + + # Zero-out Jacobian corresponding to z translation + if ndim == 3 and not self.zTranslation: + imageGradients[2][:] = 0.0 + # Coordinate frame is somehow flipped?? + + return tvol, imageGradients + dangles + + def asMatrix(self): + #XXX: grid coordinate frame is flipped + return transformationMatrix(-self.shift, -self.rotation) + + def __repr__(self): + return "EuclideanTransformation(shift=%s, rotation=%s)" % (repr(self.shift), repr(self.rotation)) + + class Displacement(Transformation, Serializable): """ @@ -92,3 +384,53 @@ def apply(self, im): def __repr__(self): return "PlanarDisplacement(delta=%s)" % repr(self.delta) + + +def transformationMatrix(shift, rot=None): + """Create an affine transformation matrix + + Parameters + ---------- + shift : array, shape (ndim,) + translations along each dimension + rot : scalar or array with shape (ndim, ) + + Returns + ------- + A : array, shape (ndims + 1, ndims + 1) + transformation matrix that shifts and rotates a set of points in homogeneous coordinates + """ + + ndim = len(shift) + if rot is None: + rot = np.zeros(ndim) + else: + rot = np.atleast_1d(rot) + c = np.cos(rot) + s = np.sin(rot) + trans = np.eye(ndim + 1) + trans[:-1, -1] = shift + xrot = np.array( + [[c[0], -s[0], 0, 0], + [s[0], c[0], 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1]]) + if ndim == 2 or np.size(rot) == 1: + A = np.dot(trans, xrot[:ndim + 1, :ndim + 1]) + else: + yrot = [[c[1], 0, s[1], 0,], + [0, 1, 0, 0], + [-s[1], 0, c[1], 0], + [0, 0, 0, 1]] + zrot = [[1, 0, 0, 0], + [0, c[2], -s[2], 0], + [0, s[2], c[2], 0], + [0, 0, 0, 1]] + A = np.dot(trans, np.dot(xrot, np.dot(yrot, zrot))) + return A + +# Dict of valid types of Transformations used by Lucas-Kanade +TRANSFORMATION_TYPES = { + 'Translation': TranslationTransformation, + 'Euclidean': EuclideanTransformation +}