Source code for ddd_pylib.image_manip._split_axis_operator

from ddd_pylib._base import ImageOperator




class SplitAxisOperator(ImageOperator):
    """
    Split an axis to get all the images on this axis (for example : time axis)

    :param axis: The axis that will be split
    """

    def __init__(self):
        super().__init__()
        self._targetAxis = None



    def setSplittingAxis(self, axis: str):
        """
        Set the axis that will be split. The axis must be present in the main image.

        Args:
            axis: Letter that define an axis
        """
        if axis != "C" and axis not in ImageOperator.validAxes:
            raise ValueError(
                f"Axis '{axis}' is not valid, must be one of {ImageOperator.validAxes} or 'C'."
            )
        self._targetAxis = axis




    def sanityCheck(self):
        super().sanityCheck()
        if self._targetAxis is None:
            raise ValueError("Splitting axis is not set.")
        if self._targetAxis not in self.getMainImage().dims:
            raise ValueError(
                f"Splitting axis '{self._targetAxis}' is not present in the main image dimensions."
            )




    def getPrefix(self):
        return f"Split[{self._targetAxis}]"


    def _run(self):
        self.sanityCheck()
        img = self.getMainImage()
        self._result = []

        for t in range(img.sizes[self._targetAxis]):
            sliced_img = img.isel({self._targetAxis: t})
            self._result.append(sliced_img)
            yield t