Source code for ddd_pylib.labeling._seeded_watershed_operator

from ddd_pylib._base import ImageOperator
from skimage.segmentation import watershed
import xarray as xr
import numpy as np
import pandas as pd




class SeededWatershedOperator(ImageOperator):
    """
    Apply seeded watershed on given image, it uses seeds, that can be calculated by FindExtrema
    if not given in the seeds attribute. The dt_image is for the FindExtrema and is a distance
    transformed image, that is calculated with DistanceTransform if not given.

    Args:
        seeds: the seeds for the seeded watershed
        dt_image: the distance transform image for the find extrema
        min_value: the minimum value of the extrema
    """
    def __init__(self):
        super().__init__()
        self._seeds = None
        self._morphological = True



    def getPrefix(self):
        return "Watershed-"




    def getSeeds(self) -> pd.DataFrame:
        """
        Get the seeds parameter
        
        Returns:
            seeds: seeds for the seeded watershed
        """
        if self._seeds is None:
            raise ValueError("Seeds are not set. Please set seeds before running the operator.")
        return self._seeds




    def setMorphological(self, morphological: bool):
        """
        Set the morphological parameter
        
        Args:
            morphological: whether to use morphological watershed or not
        """
        self._morphological = morphological




    def setSeeds(self, seeds: pd.DataFrame):
        """
        Set the seeds parameter
        
        Args:
            seeds: seeds for the seeded watershed
        """
        self._seeds = seeds


    def _seededWatershed(self, image: xr.DataArray, t: int) -> xr.DataArray:
        df = self.getSeeds()
        seeds = df if 'T' not in df.columns else df[df['T'] == t]
        coords = seeds[[str(ax) for ax in image.dims]].to_numpy()
        labels = np.arange(1, len(coords) + 1)

        seeds_mask = np.zeros(image.shape, dtype=np.uint8)
        seeds_mask[tuple(coords.T)] = labels

        kwargs = {}
        if self._morphological:
            kwargs['mask'] = image.values > 0

        result = watershed(
            image.values, 
            markers=seeds_mask, 
            **kwargs
        )

        return xr.DataArray(
            result,
            dims=image.dims,
            attrs=image.attrs
        )

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