<commit_before>"""
This module contains a set of similarity measures that was designed for use
within the Lucas-Kanade framework. They therefore expose a number of methods
that make them useful for inverse compositional and forward additive
Lucas-Kanade.

These similarity measures are designed to be dimension independent where
possible. For this reason, some methods look more complicated than would be
normally the case. For example, calculating the Hessian involves summing
a multi-dimensional array, so we dynamically calculate the list of axes
to summer over. However, the basics of the logic, other than dimension
reduction, should be similar to the original algorithms.

Citations:

Lucas, <NAME>., and <NAME>.
"An iterative image registration technique with an application to stereo
vision."
IJCAI. Vol. 81. 1981.
"""
import abc
import numpy as np
from numpy.fft import fftshift, fftn
import scipy.ndimage
import scipy.linalg
from pybug.convolution import log_gabor
from pybug.image import Image
import pybug.matlab as matlab
from pybug.warp import warp, warp_image_onto_template_image
from pybug.warp.base import map_coordinates_interpolator


class Residual(object):
    """
    An abstract base class for calculating the residual between two images
    within the Lucas-Kanade algorithm. The classes were designed
    specifically to work within the Lucas-Kanade framework and so no
    guarantee is made that calling methods on these subclasses will generate
    correct results.
    """
    __metaclass__ = abc.ABCMeta

    @property
    def error(self):
        r"""
        The RMS of the error image. Will only generate a result if the
        steepest descent update has been calculated prior.

        .. math::
            error = \sqrt{\sum_x E(x)^2}

        where :math:`E(x) = T(x) - I(W(x;p))` within the forward additive
        framework.
        """
        return np.sqrt(np.mean(self._error_img ** 2))

    @abc.abstractmethod
    def steepest_descent_images(self, image, dW_dp, **kwargs):
        r"""
        Calculates the standard steepest descent images. Within the forward
        additive framework this is defined as

        .. math::
             \nabla I \frac{\partial W}{\partial p}

        :param image: The image to calculate the steepest descent
            images from, could be either the template or input image depending
            on which framework is used.
        :type image: ndarray
        :param dW_dp: The Jacobian of the warp.
        :type dW_dp: ndarray
        :return: The steepest descent images of shape
            (num_params x image_height x image_width)
        :rtype: ndarray
        """
        pass

    @abc.abstractmethod
    def calculate_hessian(self, VT_dW_dp):
        pass

    @abc.abstractmethod
    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        pass

    def _calculate_gradients(self, image, forward=None):
        # Calculate the gradient over the image
        gradient = image.gradient()

        # Warp image for forward additive, if we've been given a transform
        if forward:
            template, transform, interp = forward
            gradient = warp_image_onto_template_image(gradient, template,
                                                      transform,
                                                      interpolator=interp)

        return gradient.as_vector(keep_channels=True)

    def _sum_over_axes(self, tensor, axes):
        tensor_summed = np.apply_over_axes(np.sum, tensor, axes)
        return np.squeeze(tensor_summed)

    def _sum_Hessian(self, H):
        # Creates a reverse list from n_dim + 1:2 to sum over
        # eg. for 3 dimensional image: [4, 3, 2]
        axes = range(len(H.shape) - 1, 1, -1)
        return self._sum_over_axes(H, axes)

    def _sum_steepest_descent(self, sd):
        # Creates a reverse list from n_dim:1 to sum over
        # eg. for 3 dimensional image: [3, 2, 1]
        axes = range(len(sd.shape) - 1, 0, -1)
        return self._sum_over_axes(sd, axes)


class LSIntensity(Residual):

    def steepest_descent_images(self, image, dW_dp, forward=None):
        gradient = self._calculate_gradients(image, forward=forward)
        return np.sum(dW_dp * gradient[:, np.newaxis, :], axis=2)

    def calculate_hessian(self, VT_dW_dp):
        return VT_dW_dp.T.dot(VT_dW_dp)

    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        self._error_img = IWxp.as_vector() - template.as_vector()
        return VT_dW_dp.T.dot(self._error_img)


class GaborFourier(Residual):

    def __init__(self, image_shape, **kwargs):
        if 'filter_bank' in kwargs:
            self._filter_bank = kwargs.get('filter_bank')
            if self._filter_bank.shape != image_shape:
                raise ValueError('Filter bank must match the '
                                 'size of the image')
        else:
            gabor = log_gabor(np.ones(image_shape), **kwargs)
            self._filter_bank = gabor[2]  # Get filter bank matrix

    def steepest_descent_images(self, image, dW_dp, **kwargs):
        gradient = self._calculate_gradients(image, dW_dp.shape[-2:],
                                             **kwargs)
        gradient = [g[np.newaxis, ...] for g in gradient]
        gradient = np.concatenate(gradient, axis=0)
        VT_dW_dp = np.sum(dW_dp * gradient[:, np.newaxis, ...], axis=0)

        # Get a range from 1 to number of image dimensions for computing
        # FFT over
        image_dims = range(1, len(image.shape) + 1)

        # Compute FFT over each parameter
        return fftshift(fftn(VT_dW_dp, axes=image_dims), axes=image_dims)

    def calculate_hessian(self, VT_dW_dp):
        filtered_jac = (self._filter_bank ** 0.5) * VT_dW_dp
        H = np.conjugate(filtered_jac[:, np.newaxis, ...]) * filtered_jac
        return self._sum_Hessian(H)

    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        self._error_img = fftshift(fftn(IWxp - template))
        ft_error_img = self._filter_bank * self._error_img
        sd = VT_dW_dp * np.conjugate(ft_error_img)
        return self._sum_steepest_descent(sd)


class ECC(Residual):

    def __normalise_images(self, image):
        # TODO: do we need to copy the image?
        new_im = Image(image.pixels, mask=image.mask)
        i = new_im.pixels
        i -= np.mean(i)
        i /= scipy.linalg.norm(i)
        i = np.nan_to_num(i)

        new_im.pixels = i
        return new_im

    def steepest_descent_images(self, image, dW_dp, forward=None):
        norm_image = self.__normalise_images(image)

        gradient = self._calculate_gradients(norm_image, forward=forward)
        return np.sum(dW_dp * gradient[:, np.newaxis, :], axis=2)

    def calculate_hessian(self, VT_dW_dp):
        H = VT_dW_dp.T.dot(VT_dW_dp)
        self.__H_inv = scipy.linalg.inv(H)

        return H

    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        normalised_IWxp = self.__normalise_images(IWxp).as_vector()
        normalised_template = self.__normalise_images(template).as_vector()

        Gt = VT_dW_dp.T.dot(normalised_template)
        Gw = VT_dW_dp.T.dot(normalised_IWxp)

        # Calculate the numerator
        IWxp_norm = scipy.linalg.norm(normalised_IWxp)
        num = (IWxp_norm ** 2) - np.dot(Gw.T, np.dot(self.__H_inv, Gw))

        # Calculate the denominator
        den1 = np.dot(normalised_template, normalised_IWxp)
        den2 = np.dot(Gt.T, np.dot(self.__H_inv, Gw))
        den = den1 - den2<commit_msg>Change _calculate_gradients to return an image<commit_after>"""
This module contains a set of similarity measures that was designed for use
within the Lucas-Kanade framework. They therefore expose a number of methods
that make them useful for inverse compositional and forward additive
Lucas-Kanade.

These similarity measures are designed to be dimension independent where
possible. For this reason, some methods look more complicated than would be
normally the case. For example, calculating the Hessian involves summing
a multi-dimensional array, so we dynamically calculate the list of axes
to summer over. However, the basics of the logic, other than dimension
reduction, should be similar to the original algorithms.

Citations:

Lucas, <NAME>., and <NAME>.
"An iterative image registration technique with an application to stereo
vision."
IJCAI. Vol. 81. 1981.
"""
import abc
import numpy as np
from numpy.fft import fftshift, fftn
import scipy.ndimage
import scipy.linalg
from pybug.convolution import log_gabor
from pybug.image import Image
import pybug.matlab as matlab
from pybug.warp import warp, warp_image_onto_template_image
from pybug.warp.base import map_coordinates_interpolator


class Residual(object):
    """
    An abstract base class for calculating the residual between two images
    within the Lucas-Kanade algorithm. The classes were designed
    specifically to work within the Lucas-Kanade framework and so no
    guarantee is made that calling methods on these subclasses will generate
    correct results.
    """
    __metaclass__ = abc.ABCMeta

    @property
    def error(self):
        r"""
        The RMS of the error image. Will only generate a result if the
        steepest descent update has been calculated prior.

        .. math::
            error = \sqrt{\sum_x E(x)^2}

        where :math:`E(x) = T(x) - I(W(x;p))` within the forward additive
        framework.
        """
        return np.sqrt(np.mean(self._error_img ** 2))

    @abc.abstractmethod
    def steepest_descent_images(self, image, dW_dp, **kwargs):
        r"""
        Calculates the standard steepest descent images. Within the forward
        additive framework this is defined as

        .. math::
             \nabla I \frac{\partial W}{\partial p}

        :param image: The image to calculate the steepest descent
            images from, could be either the template or input image depending
            on which framework is used.
        :type image: ndarray
        :param dW_dp: The Jacobian of the warp.
        :type dW_dp: ndarray
        :return: The steepest descent images of shape
            (num_params x image_height x image_width)
        :rtype: ndarray
        """
        pass

    @abc.abstractmethod
    def calculate_hessian(self, VT_dW_dp):
        pass

    @abc.abstractmethod
    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        pass

    def _calculate_gradients(self, image, forward=None):
        # Calculate the gradient over the image
        gradient = image.gradient()

        # Warp image for forward additive, if we've been given a transform
        if forward:
            template, transform, interp = forward
            gradient = warp_image_onto_template_image(gradient, template,
                                                      transform,
                                                      interpolator=interp)

        return gradient


class LSIntensity(Residual):

    def steepest_descent_images(self, image, dW_dp, forward=None):
        gradient = self._calculate_gradients(image, forward=forward)
        gradient = gradient.as_vector(keep_channels=True)
        return np.sum(dW_dp * gradient[:, np.newaxis, :], axis=2)

    def calculate_hessian(self, VT_dW_dp):
        return VT_dW_dp.T.dot(VT_dW_dp)

    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        self._error_img = IWxp.as_vector() - template.as_vector()
        return VT_dW_dp.T.dot(self._error_img)


class GaborFourier(Residual):

    def __init__(self, image_shape, **kwargs):
        if 'filter_bank' in kwargs:
            self._filter_bank = kwargs.get('filter_bank')
            if self._filter_bank.shape != image_shape:
                raise ValueError('Filter bank must match the '
                                 'size of the image')
        else:
            gabor = log_gabor(np.ones(image_shape), **kwargs)
            self._filter_bank = gabor[2]  # Get filter bank matrix

    def steepest_descent_images(self, image, dW_dp, **kwargs):
        gradient = self._calculate_gradients(image, dW_dp.shape[-2:],
                                             **kwargs)
        gradient = [g[np.newaxis, ...] for g in gradient]
        gradient = np.concatenate(gradient, axis=0)
        VT_dW_dp = np.sum(dW_dp * gradient[:, np.newaxis, ...], axis=0)

        # Get a range from 1 to number of image dimensions for computing
        # FFT over
        image_dims = range(1, len(image.shape) + 1)

        # Compute FFT over each parameter
        return fftshift(fftn(VT_dW_dp, axes=image_dims), axes=image_dims)

    def calculate_hessian(self, VT_dW_dp):
        filtered_jac = (self._filter_bank ** 0.5) * VT_dW_dp
        H = np.conjugate(filtered_jac[:, np.newaxis, ...]) * filtered_jac
        return self._sum_Hessian(H)

    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        self._error_img = fftshift(fftn(IWxp - template))
        ft_error_img = self._filter_bank * self._error_img
        sd = VT_dW_dp * np.conjugate(ft_error_img)
        return self._sum_steepest_descent(sd)


class ECC(Residual):

    def __normalise_images(self, image):
        # TODO: do we need to copy the image?
        new_im = Image(image.pixels, mask=image.mask)
        i = new_im.pixels
        i -= np.mean(i)
        i /= scipy.linalg.norm(i)
        i = np.nan_to_num(i)

        new_im.pixels = i
        return new_im

    def steepest_descent_images(self, image, dW_dp, forward=None):
        norm_image = self.__normalise_images(image)

        gradient = self._calculate_gradients(norm_image,
                                             forward=forward)
        gradient = gradient.as_vector(keep_channels=True)
        return np.sum(dW_dp * gradient[:, np.newaxis, :], axis=2)

    def calculate_hessian(self, VT_dW_dp):
        H = VT_dW_dp.T.dot(VT_dW_dp)
        self.__H_inv = scipy.linalg.inv(H)

        return H

    def steepest_descent_update(self, VT_dW_dp, IWxp, template):
        normalised_IWxp = self.__normalise_images(IWxp).as_vector()
        normalised_template = self.__normalise_images(template).as_vector()

        Gt = VT_dW_dp.T.dot(normalised_template)
        Gw = VT_dW_dp.T.dot(normalised_IWxp)

        # Calculate the numerator
        IWxp_norm = scipy.linalg.norm(normalised_IWxp)
        num = (IWxp_norm ** 2) - np.dot(Gw.T, np.dot(self.__H_inv, Gw))

        # Calculate the denominator
        den1 = np.dot(normalised_template, normalised_IWxp)
        den2 = np.dot(Gt.T, np.dot(self.__H_inv, Gw))
        den = den1 - den2
{'fraction_non_alphanumeric': 0.08523154798549833, 'fraction_numerical': 0.004446268554620699, 'mean_word_length': 2.582455280568488, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '110633', 'n_tokens_mistral': 4621, 'n_tokens_neox': 4428, 'n_words': 1281}