Source code for ddd_pylib.image_manip._reslice_operator

from ddd_pylib._base import ImageOperator
import numpy as np




class ResliceOperator(ImageOperator):
    """
    Change the axis through which the image is viewed.
    The direction can be set to any of the following: +T, -T, +Z, -Z, +Y, -Y, +X, -X.
    """

    _DIRECTIONS = set(["+T", "-T", "+Z", "-Z", "+Y", "-Y", "+X", "-X"])

    def __init__(self):
        super().__init__()
        self._direction = self.defaultDirection()
        self._keepTimePosition = self.defaultKeepTimePosition()



    @staticmethod
    def defaultDirection():
        return "+Z"




    @staticmethod
    def defaultKeepTimePosition():
        return True




    def setDirection(self, direction: str):
        """
        Set the new axis to be viewed along

        Args:
            direction: The axis through which the image should be resliced.
                       Add '-' in front to flip the image along that axis.
        """
        if direction not in self._DIRECTIONS:
            raise ValueError(f"The direction '{direction}' is not valid.")
        self._direction = direction




    def getPrefix(self):
        return "Resliced-"




    def sanityCheck(self):
        super().sanityCheck()
        img = self.getMainImage()
        if "Z" not in img.dims:
            raise ValueError("2D images aren't supported")


    def _rollTo(self, dims, ax):
        i = dims.index(ax)
        rolled = np.roll(dims, -i).tolist()
        if "T" not in rolled:
            return rolled
        t_pos = dims.index("T")
        if self._keepTimePosition:
            rolled.remove("T")
            rolled.insert(t_pos, "T")
        return rolled



    def getOutputDims(self):
        im = self.getMainImage()
        ax = self._direction[1]
        return self._rollTo(list(im.dims), ax)


    def _applyToFrame(self, frameData, t, nT):
        d = self._direction
        return (
            frameData[0]
            if d[0] != "-"  # Flip axis on the fly
            else frameData[0].isel({d[1]: slice(None, None, -1)})
        )