Source code for ddd_pylib.image_manip._euclidean_distance_transform_operator

from ddd_pylib._base import ImageOperator
from ddd_pylib import ImageUtils
from scipy.ndimage import distance_transform_edt
import xarray as xr
import numpy as np




class EuclideanDistanceTransformOperator(ImageOperator):
    """
    Creates the Euclidean distance transform (EDT) of the main image.
    The anisotropy of the image is taken into account by using the scale of each axis as a sampling parameter.
    The EDT assigns to each pixel the distance to the nearest zero pixel.
    The provided input image must be a binary mask, with only two unique values.
    """

    def __init__(self):
        super().__init__()



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




    def sanityCheck(self):
        """
        Checks that the main image is a binary mask, with only two unique values.
        """
        super().sanityCheck()
        img = self.getMainImage()
        unique = np.unique(img.values)
        if len(unique) > 2:
            raise ValueError("Image must be binary for distance transform. Found more than two unique values.")


    def _distanceTransform(self, image: xr.DataArray) -> xr.DataArray:
        scales = ImageUtils.scaleToTuple(image)
        res = distance_transform_edt(
            image.values, 
            sampling=scales
        )
        return xr.DataArray(
            res, 
            dims=image.dims, 
            attrs=image.attrs
        )

    def _applyToFrame(self, frameData, t, nT):
        image = frameData[0]
        return self._distanceTransform(image)