Source code for ddd_pylib.labeling._threshold_operator

from ddd_pylib._base import ImageOperator
from skimage import filters
from typing import Tuple
import xarray as xr
import numpy as np




class ThresholdOperator(ImageOperator):
    """
    Creates a thresholded image (binary mask). 
    The thresholding method can be selected from:

    - Isodata
    - Li
    - Mean
    - Median
    - Minimum
    - Otsu
    - Triangle
    - Yen
    - Local
    - Niblack
    - Sauvola

    The default being **Otsu**.
    
    This operator can handle objects on light and dark background. If you set the blackBackground to True, 
    the higher values will be the foreground (True in the resulting mask). If you set it to False, the 
    lower values will be the foreground.

    The threshold can be computed independently for each time point, or globally for the whole image.

    If you set the threshold manually, it will override the thresholding method. 
    This range will be the same for all time points whether you activated the independentTime option or not.

    Settable parameters:
        thresholdMethod (setThresholdMethod): The thresholding method to use
        independentTime (setIndependentTime): If True, the threshold is computed independently for each time point. If False, the threshold is computed globally for the whole image.
        blackBackground (setBlackBackground): If True, the higher values will be the foreground. If False, the lower values will be the foreground.
        thresholdValue (setThresholdValue): If set, the threshold is fixed to this value, and the thresholdMethod is ignored.
        thresholdRange (setThresholdRange): If set, the threshold is fixed to this range, and the thresholdMethod is ignored. The range is defined by a low and a high value. The foreground will be the values within this range.

    Produces:
        result (getResult): The thresholded image as a :code:`xr.DataArray` with the same dimensions and attributes as the main image.
                            Values are boolean, with True for foreground and False for background.

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

    methods = {
        'Isodata'  : filters.threshold_isodata,
        'Li'       : filters.threshold_li,
        'Mean'     : np.mean,
        'Median'   : np.median,
        'Minimum'  : filters.threshold_minimum,
        'Otsu'     : filters.threshold_otsu,
        'Triangle' : filters.threshold_triangle,
        'Yen'      : filters.threshold_yen,
        'Local'    : filters.threshold_local,
        'Niblack'  : filters.threshold_niblack,
        'Sauvola'  : filters.threshold_sauvola
    }

    def __init__(self):
        super().__init__()
        self._method = self.defaultMethod()
        self._blackBackground = self.defaultBlackBackground()
        self._independentTime = self.defaultIndependentTime()
        self._fixedRange = self.defaultFixedRange()



    @staticmethod
    def defaultMethod() -> str:
        return "Otsu"




    @staticmethod
    def defaultBlackBackground() -> bool:
        return True




    @staticmethod
    def defaultIndependentTime() -> bool:
        return False




    @staticmethod
    def defaultFixedRange() -> Tuple:
        return (-np.inf, np.inf)




    def getPrefix(self):
        return f"Threshold[{self._method}]-"




    def setThresholdMethod(self, method):
        """
        Take a method name and set it in the method attribute 

        Args:
            method: A string defining a thresholding method name, which needs to be in the methods set
        """
        if method not in self.methods:
            raise ValueError("Unknown method")
        self._method = method




    def setIndependentTime(self, independentTime: bool):
        """
        Set the value of the independentTime attribute
        """
        self._independentTime = independentTime




    def setBlackBackground(self, blackBackground: bool):
        """
        Set the value of the blackBackground attribute
        """
        self._blackBackground = blackBackground




    def setThresholdValue(self, value):
        self._fixedRange = (
            (value, np.inf) 
            if self._blackBackground 
            else (-np.inf, value)
        )
        self._method = None




    def setThresholdRange(self, low, high):
        if low > high:
            raise ValueError("Low value can't be higher than high value")
        self._fixedRange = (low, high)
        self._method = None


    def _processThreshold(self, image: xr.DataArray) -> Tuple:
        if self._method is None:
            return (
                self._fixedRange 
                if self._fixedRange is not None 
                else (-np.inf, np.inf)
            )
        else:
            threshold_func = self.methods[self._method]
            threshold_value = threshold_func(image.values)
            return (
                (threshold_value, np.inf) 
                if self._blackBackground 
                else (-np.inf, threshold_value)
            )

    def _threshold(self, image: xr.DataArray):
        bounds = (
            self._processThreshold(image) 
            if self._independentTime 
            else self._fixedRange
        )

        if bounds is None:
            raise ValueError("Threshold bounds are not set.")

        low, high = bounds
        thresh_image = None

        if self._blackBackground:
            thresh_image = (
                (image.data >= low) &
                (image.data <= high)
            )
        else:
            thresh_image = (
                (image.data <= low) &
                (image.data >= high)
            )

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

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

    def _run(self):
        if not self._independentTime:
            img = self.getMainImage()
            self._fixedRange = self._processThreshold(img)
        yield from super()._run()