Source code for ddd_pylib.image_manip._crop_operator

from ddd_pylib._base import ImageOperator
import xarray as xr
from typing import List




class CropOperator(ImageOperator):
    """
    Takes a list of requested crops and extracts them from the main image. Each crop is a dictionary with:

    - **keys (str):** The axes to crop (e.g. "X", "Y", "Z", "T")
    - **values (tuple[int, int]):** The start and end points of the crop along that axis

    The tuples follow the same logic as Python slicing, where the start is inclusive and the end is exclusive.
    For example, a crop of (0, 100) along the X axis will include pixels from 0 to 99.

    Example of a crop list:

    .. code-block:: python

        [
            {"X": (0, 100), "Y": (20, 200), "Z": (0, 10), "T": (2, 3)},
            {"X": (100, 200), "Y": (0, 200), "Z": (0, 10)}
        ]

    If a crop dict is not mentioning an axis, the whole axis is taken (for example, if 'T' is missing,
    the whole time series will be exported). A crop with an empty dict will take the whole image.
    The result is a list of tuples with:

    - the first element of the tuples represents the translation of the cropped image in the original image (the starting point of the crop)
    - the second element of the tuples is the cropped image as an `xr.DataArray`

    Settable parameters:
        mainImage (setMainImage): The main image to crop
        crops (setCropsList): The list of crops to apply to the image

    Produces:
        result (getResult): A list of tuples with the translation and the cropped image

    Example:
        .. literalinclude:: /_examples/crop_operator.py
           :language: python
    """

    def __init__(self):
        super().__init__()
        self._userCrops = []

    def _isCropValid(self, crop: List[dict]):
        if len(crop) == 0:
            raise ValueError("At least one crop is required")
        img = self.getMainImage()
        for user_crop in crop:
            for axis, (start, end) in user_crop.items():
                axis = str(axis)
                if start >= end:
                    raise ValueError(
                        f"End point can't be behind start point : {(start, end)}"
                    )
                if start < 0:
                    raise ValueError(f"Start point can't be negative : {start}")
                if end > img.sizes.get(axis, 0):
                    raise ValueError(
                        f"End point can't be bigger than the image size : {end} > {img.sizes.get(axis, 0)}"
                    )



    def getPrefix(self) -> str:
        return "Cropped-"




    def sanityCheck(self):
        """
        Verifies that the requested crops are valid:

        - The start point must be less than the end point
        - The start point must be non-negative
        - The end point must be less than or equal to the size of the image along that axis
        """
        super().sanityCheck()
        self._isCropValid(self._userCrops)




    def setCropsList(self, crops: List[dict]):
        """
        Sets the list of crops that is needed to get from the image.

        Args:
            crops: List of dictionaries that define crops.
        """
        self._userCrops = crops


    def _addMissingAxes(self, crop):
        img = self.getMainImage()
        crop = crop.copy()
        for ax, size in img.sizes.items():
            if ax not in crop:
                crop[str(ax)] = (0, size)
        return crop

    def _applyToFrame(self, frameData, t, nT):
        raise RuntimeError("Crop operator is not meant to be applied framewise")

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

        for t, crop in enumerate(self._userCrops):
            crop = self._addMissingAxes(crop)
            decapsulated = [crop[str(ax)] for ax in self.getMainImage().dims]
            crop_img = img.values[
                tuple(slice(start, end) for start, end in decapsulated)
            ]
            self._result.append(
                (
                    tuple(crop[str(ax)][0] for ax in img.dims),
                    xr.DataArray(
                        crop_img,
                        dims=img.dims,
                        attrs=img.attrs,
                        name=self.makeResultName(t),
                    ),
                )
            )
            yield t