Source code for ddd_pylib._base._image_operator

from ._base_operator import BaseOperator
from ddd_pylib import ImageUtils
import xarray as xr
from typing import List, Tuple
from abc import abstractmethod




class ImageOperator(BaseOperator):
    """
    Gives all the common functions that all Operator that use images will use.
    It provides an image attribute which stock the main image used and a 
    result attribute which stock the result of the function used.

    :param _mainImage: Take a DataArray defining the image
    :param _result: Store the result of the function used, its form may vary depending of the function
    """

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


    def getMainImage(self) -> xr.DataArray:
        """
        Return the main image as a DataArray
        """
        if self._mainImage is None:
            raise ValueError("Main image is not set")
        return self._mainImage




    def setMainImage(self, image: xr.DataArray):
        """
        Take an image and store it in the _mainImage attribute

        Args:
            image: An image as a DataArray
        """
        self._mainImage = image




    def makeResultName(self, t=None) -> str:
        """
        Create the result name using the image name and the getPrefix function of the used operator
        """
        if self._mainImage is None:
            raise ValueError("Main image is not set")
        name = f"{self.getPrefix()}{self._mainImage.name if self._mainImage.name is not None else ''}"
        if t is not None:
            name = f"{name}[{t+1}]"
        return name




    def sanityCheck(self):
        """
        Simple check to prevent simple error, it may be implemented in subclasses
        to check more precise stuff
        """
        if self._mainImage is None:
            raise ValueError("Main image is not set")
        currentAxes = set([str(a) for a in self._mainImage.dims])
        if not currentAxes.issubset(self.getValidAxes()):
            raise ValueError(f"Main image dimensions {self._mainImage.dims} are not valid. Valid axes are {self.getValidAxes()}")


    @abstractmethod
    def _applyToFrame(self, frameData: Tuple[xr.DataArray, ...], t: int, nT: int):
        raise NotImplementedError("Subclasses should implement this method")



    def getNTimePoints(self) -> int:
        """
        Return the size of the T axis if it exist, if not then just 1
        """
        img = self.getMainImage()
        return img.sizes.get('T', 1)


    def _zipInputs(self) -> List[Tuple[xr.DataArray, ...]]:
        img = self.getMainImage()
        if 'T' in img.dims:
            return [(img.isel(T=t),) for t in range(self.getNTimePoints())]
        else:
            return [(img,)]



    def getOutputDims(self) -> List[str]:
        img = self.getMainImage()
        return [str(a) for a in img.dims]


    def _bundleResultsBuffer(self, buffer: List):
        """
        Bundle the results in the buffer into a single DataArray or DataFrame.
        If the main image has a T dimension, concatenate along that dimension.
        Otherwise, return the first element of the buffer.
        """
        img = self.getMainImage()
        result = (
            xr.concat(buffer, dim='T').transpose(*self.getOutputDims())
            if 'T' in img.dims
            else buffer[0]
        )
        result.name = self.makeResultName()
        ImageUtils.completeCalibrationFrom(img, result)
        return result