Source code for ddd_pylib.image_manip._normalize_operator

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




class NormalizeOperator(ImageOperator):
    """
    Take an image and normalizes it.
    If the image has a time axis (T), it is possible with setIndividualTimePoints
    to normalize them individually or to use the global range of values.
    It is also possible to set a range in which the image should be normalized with setTargetRange.
    """

    def __init__(self):
        super().__init__()
        self._individualTimePoints = self.defaultIndividualTimePoints()
        self._targetRange = self.defaultTargetRange()
        self._extremas = None



    @staticmethod
    def defaultIndividualTimePoints():
        return False




    @staticmethod
    def defaultTargetRange():
        return None




    def setIndividualTimePoints(self, individual: bool):
        """
        Set whether to normalize individual time points or use the global range.
        """
        self._individualTimePoints = individual




    def setTargetRange(self, target_range: tuple):
        """
        Set the target range in which the images should be normalized.

        Args:
            target_range: The target range as a tuple.
        """
        start, end = target_range
        if start >= end:
            raise ValueError("Range length can't be 0 or under")
        if start < 0:
            raise ValueError("Range start can't be negative")
        self._targetRange = target_range




    def getTargetRange(self) -> Tuple:
        """
        Return the target range in which the images should be normalized.
        """
        if self._targetRange is not None:
            return self._targetRange
        else:
            dt = self.getMainImage().dtype
            if np.issubdtype(dt, np.floating):
                return (0.0, 1.0)
            if np.issubdtype(dt, np.integer):
                info = np.iinfo(dt)
                return (info.min, info.max)
        raise ValueError("Cannot determine target range for the image type.")


    def _processExtremas(self):
        img = self.getMainImage()
        if self._individualTimePoints:
            self._extremas = [
                (img.isel(T=t).values.min(), img.isel(T=t).values.max())
                for t in range(self.getNTimePoints())
            ]
        else:
            self._extremas = [
                (img.values.min(), img.values.max())
            ] * self.getNTimePoints()

    def _normalize(
        self, image: xr.DataArray, extremas: Tuple[int, int]
    ) -> xr.DataArray:
        lower, upper = extremas
        vals = image.values.astype(np.float32)
        vals = (vals - lower) / (upper - lower)
        start, end = self.getTargetRange()
        vals = vals * (end - start) + start
        return xr.DataArray(
            data=vals, dims=image.dims, attrs=image.attrs.copy()
        ).astype(image.dtype)

    def _applyToFrame(self, frameData, t, nT):
        if self._extremas is None:
            raise ValueError(
                "Extremas have not been processed. Call _processExtremas() before applying normalization."
            )
        return self._normalize(frameData[0], self._extremas[t])

    def _run(self):
        self._processExtremas()
        yield from super()._run()



    def getPrefix(self):
        return "Normalized-"