Source code for ddd_pylib.labels_manip._replace_values_operator

from ddd_pylib._base import ImageOperator, MeasurementsOperator
import numpy as np
import xarray as xr
from typing import List, Dict, Tuple
from skimage.segmentation import clear_border
import pandas as pd




class ReplaceValuesOperator(ImageOperator):
    """
    Takes an image and a dictionary of values representing a LUT and replaces the values in 
    the image according to this LUT. Only images with integer dtypes are supported.
    """
    def __init__(self):
        super().__init__()
        self._valuesDict = None
        self._lut = None



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

    


    def setValuesDict(self, values_dict: Dict[np.integer, np.integer]|Dict[int, int]):
        """
        Set the dictionary of values to be replaced. The keys are the original values, and the values are the new values.
        If a value is not specified in the dictionary, it will be replaced by itself (i.e., it will remain unchanged).
        """
        self._valuesDict = values_dict




    def sanityCheck(self):
        super().sanityCheck()
        dt = self.getMainImage().dtype
        if not np.issubdtype(dt, np.integer):
            raise TypeError("Main image must have an integer dtype.")


    def _createLUT(self) -> np.ndarray:
        self.sanityCheck()
        if self._valuesDict is None:
            raise ValueError("Values dictionary is not set.")
        img = self.getMainImage()
        # size of the LUT
        max_value = int(np.max(img.values)) + 1
        as_vector = [(k, v) for k, v in self._valuesDict.items()]
        as_vector = sorted(as_vector, key=lambda x: x[0])
        # list of original values
        orig = np.array([k for k, _ in as_vector])
        # what each original value should be mapped to
        mapped_values = np.array([v for _, v in as_vector])
        # LUT: indices = original values, values = new values
        # default: values mapped to themselves
        lut = np.arange(max_value, dtype=mapped_values.dtype)
        # filling the LUT with the new value values
        lut[orig] = mapped_values
        return lut

    def _replaceValues(self, image: xr.DataArray) -> xr.DataArray:
        if self._lut is None:
            raise ValueError("LUT is not created. Call _run() before applying value replacement.")
        new_image = image.copy()
        # broadcasting the shape of 'image' through the LUT to get the new values map
        new_image.values = self._lut[image.values]
        return new_image

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

    def _run(self):
        self._lut = self._createLUT()
        yield from super()._run()