Source code for ewoksndreg.transformation.scikitimage_backend

from typing import Dict
from typing import Optional
from typing import Sequence

import numpy
from skimage.transform import AffineTransform
from skimage.transform import EuclideanTransform
from skimage.transform import ProjectiveTransform
from skimage.transform import SimilarityTransform
from skimage.transform import warp

from .base import Transformation
from .homography import Homography
from .homography import reverse_indices
from .lstsq import calc_translation




class ShiftTransform(ProjectiveTransform):
    # Inspired by the implementation of EuclideanTransform
    def __init__(
        self,
        matrix: Optional[numpy.ndarray] = None,
        translation: Optional[numpy.ndarray] = None,
        dimensionality: Optional[int] = None,
    ) -> None:
        if matrix is not None and translation is not None:
            raise ValueError(
                "You cannot specify both a transformation matrix and a translation."
            )

        if matrix is not None:
            if matrix.shape[0] != matrix.shape[1]:
                raise ValueError("Invalid shape of transformation matrix.")
            if dimensionality is None:
                dimensionality = matrix.shape[1] - 1
            final_matrix = matrix
        elif translation is not None:
            translation = numpy.asarray(translation)
            if dimensionality is None:
                dimensionality = translation.size
            elif translation.size != dimensionality:
                raise ValueError(
                    f"Translation must have {dimensionality} elements for {dimensionality}D."
                )
            final_matrix = numpy.eye(dimensionality + 1)
            final_matrix[0:dimensionality, dimensionality] = translation
        else:
            if dimensionality is None:
                dimensionality = 2
            final_matrix = numpy.eye(dimensionality + 1)

        if final_matrix.shape[1] != dimensionality + 1:
            raise ValueError(
                f"Invalid matrix shape {final_matrix.shape}. "
                f"Expected a {(dimensionality + 1, dimensionality + 1)} matrix for {dimensionality}D."
            )

        super().__init__(matrix=matrix, dimensionality=dimensionality)



    @classmethod
    def from_estimate(cls, src: numpy.ndarray, dst: numpy.ndarray) -> "ShiftTransform":
        """New since scikit-image 0.26 to replace `estimate`."""
        return cls(matrix=_translation_estimate(src, dst))


    def _estimate(self, src: numpy.ndarray, dst: numpy.ndarray) -> None:
        """New since scikit-image 0.26 to replace `estimate`."""
        self.params[...] = _translation_estimate(src, dst)



    def estimate(self, src: numpy.ndarray, dst: numpy.ndarray) -> bool:
        """Deprecated since scikit-image 0.26 in favor of `from_estimate` and `_estimate`."""
        self.params[...] = _translation_estimate(src, dst)
        return True


    @property
    def translation(self) -> numpy.ndarray:
        return self.params[0 : self.dimensionality, self.dimensionality]



def _translation_estimate(src: numpy.ndarray, dst: numpy.ndarray) -> numpy.ndarray:
    src = numpy.asarray(src).T
    dst = numpy.asarray(dst).T
    return calc_translation(src, dst)




class SciKitImageHomography(
    Homography, registry_id=Homography.RegistryId("Homography", "SciKitImage")
):
    def __init__(self, *args, warp_options: Optional[Dict] = None, **kw) -> None:
        if warp_options is None:
            warp_options = dict()
        self._warp_options = warp_options

        super().__init__(*args, **kw)

        self._switched_passive = reverse_indices(self.passive_matrix)
        if self._transfo_type == self._transfo_type.identity:
            self._sc_passive = EuclideanTransform(matrix=self._switched_passive)
        elif self._transfo_type == self._transfo_type.translation:
            self._sc_passive = EuclideanTransform(matrix=self._switched_passive)
        elif self._transfo_type == self._transfo_type.rigid:
            self._sc_passive = EuclideanTransform(matrix=self._switched_passive)
        elif self._transfo_type == self._transfo_type.similarity:
            self._sc_passive = SimilarityTransform(matrix=self._switched_passive)
        elif self._transfo_type == self._transfo_type.affine:
            self._sc_passive = AffineTransform(matrix=self._switched_passive)
        elif self._transfo_type == self._transfo_type.projective:
            self._sc_passive = ProjectiveTransform(matrix=self._switched_passive)
        else:
            raise ValueError(f"'{self._transfo_type}' not supported")
        self._sc_active = self._sc_passive.inverse



    def apply_coordinates(self, coord: Sequence[numpy.ndarray]) -> numpy.ndarray:
        """
        :param coord: shape `(N, M)`
        :returns: shape `(N, M)`
        """
        res = self._sc_active(numpy.transpose(coord[::-1]))
        return numpy.transpose(res)[::-1]




    def apply_data(
        self,
        data: numpy.ndarray,
        offset: Optional[numpy.ndarray] = None,
        shape: Optional[numpy.ndarray] = None,
        cval=numpy.nan,
        interpolation_order: int = 1,
    ) -> numpy.ndarray:
        """
        :param data: shape `(N1, N2, ..., M1, M2, ...)` with `len((N1, N2, ...)) = N`
        :param offset: shape `(N,)`
        :param shape: shape `(N,) = [N1', N2', ...]`
        :param cval: missing value
        :param interpolation_order: order of interpolation: 0 is nearest neighbor, 1 is bilinear,...
        :returns: shape `(N1', N2', ..., M1, M2, ...)`
        """
        kw = dict(self._warp_options)
        if shape is not None:
            kw["output_shape"] = shape
        if offset is not None:
            kw["offset"] = offset
        if cval is not None:
            kw["cval"] = cval
        if interpolation_order in [0, 1, 3]:
            kw["order"] = interpolation_order
        else:
            raise ValueError(
                "This type of interpolation is not supported with ScikitImage"
            )
        # strange behaviour for ints with higher interpolation order
        if numpy.issubdtype(data.dtype, numpy.integer):
            data = data.astype(numpy.float64)
        # TODO: offset
        return warp(data, self._sc_passive, **kw)


    def __matmul__(self, other: Transformation):
        if isinstance(other, SciKitImageHomography):
            if self.passive_matrix.shape == other.passive_matrix.shape:
                return SciKitImageHomography(
                    other.passive_matrix @ self.passive_matrix,
                    warp_options=self._warp_options,
                )
            else:
                raise TypeError("Homographies must have same dimensions")
        else:
            raise ValueError("Only concatenation of same types allowed")