analyseScript/Analyser/ImagingAnalyser.py

import numpy as np
import xarray as xr
import copy

from DataContainer.ReadData import read_hdf5_file
from Analyser.FringeRemoval import FringeRemoval
from ToolFunction.ToolFunction import get_scanAxis


class ImageAnalyser():
    """A class for operate with and analyse images
    """

    def __init__(self) -> None:
        """Initialize the class
        """
        self._image_name = {
            'atoms': 'atoms',
            'background': 'background',
            'dark': 'dark',
            'OD':'OD',
            'optimumBackground':'optimumBackground'
        }
        self._center = None
        self._span = None
        self._fraction = None

        self._fringeRemoval = FringeRemoval()
        self.fringeRemovalReferenceImages = None
        
    @property
    def image_name(self):
        """The getter of the names of three standard images for absorption images

        :return: The names of three standard images for absorption images 
        :rtype: dict
        """
        return self._image_name
    
    @image_name.setter
    def image_name(self, value):
        """The setter of the names of three standard images for absorption images.
        It has to be a dict and have the follow format:
            {
                'atoms': the name of the image with atoms,
                'background': the name of the image without atoms,
                'dark': the name of the image without light,
                'OD': the name of the calculated OD,
            }
        
        :param value: The names of three standard images for absorption images.
        :type value: dict
        """
        self._image_name.update(value)

    @property
    def center(self):
        """The getter of the center of region of insterest (ROI)

        :return: The center of region of insterest (ROI)
        :rtype: tuple
        """
        return self._center
    
    @center.setter
    def center(self, value):
        """The setter of the center of region of insterest (ROI)

        :param value: The center of region of insterest (ROI)
        :type value: tuple
        """
        self._center = value

    @property
    def span(self):
        """The getter of the span of region of insterest (ROI) 

        :return: The span of region of insterest (ROI) 
        :rtype: tuple
        """
        return self._span

    @span.setter
    def span(self, value):
        """The setter of the span of region of insterest (ROI) 

        :param value: The span of region of insterest (ROI) 
        :type value: tuple
        """
        self._span = value

    @property
    def fraction(self):
        """The getter of the fraction used to remove the background in the region of insterest (ROI)

        :return: The fraction used to remove the background
        :rtype: tuple
        """
        return self._fraction

    @fraction.setter
    def fraction(self, value):
        """The setter of the fraction used to remove the background in the region of insterest (ROI)

        :param value: The fraction used to remove the background
        :type value: tuple
        """
        self._fraction = value
    
    def get_offset_from_corner(self, dataArray, x_fraction=None, y_fraction=None, fraction=None, xAxisName='x', yAxisName='y'):
        """Calculate the background value based on the four coners of the image.

        :param dataArray: The image
        :type dataArray: xarray DataArray
        :param x_fraction: The fraction of the pixels used x in axis, defaults to None
        :type x_fraction: float, optional
        :param y_fraction: The fraction of the pixels used y in axis, defaults to None
        :type y_fraction: float, optional
        :param fraction: The fraction of the pixels used to calculate the background, defaults to None
        :type fraction: tuple, optional
        :param xAxisName: The name of the x axis in dataArray, defaults to 'x'
        :type xAxisName: str, optional
        :param yAxisName: The name of the y axis in dataArray, defaults to 'y'
        :type yAxisName: str, optional
        :return: The value of the background
        :rtype: float
        """
        
        if fraction is None:
            if x_fraction is None:
                x_fraction = self._fraction[0]
        
            if y_fraction is None:
                y_fraction = self._fraction[1]
        else:
            x_fraction = fraction[0]
            y_fraction = fraction[1]

        x_number = dataArray[xAxisName].shape[0]
        y_number = dataArray[yAxisName].shape[0]

        mean = dataArray.isel(x=slice(0, int(x_number * x_fraction)), y=slice(0 , int(y_number * y_fraction))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(0, int(x_number * x_fraction)), y=slice(int(y_number - y_number * y_fraction) , int(y_number))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(int(x_number - x_number * x_fraction) , int(x_number)), y=slice(0 , int(y_number * y_fraction))).mean(dim=[xAxisName, yAxisName])
        mean += dataArray.isel(x=slice(int(x_number - x_number * x_fraction) , int(x_number)), y=slice(int(y_number - y_number * y_fraction) , int(y_number))).mean(dim=[xAxisName, yAxisName])
    
        return mean / 4
    
    def substract_offset(self, dataArray, **kwargs):
        """Remove the backgournd 

        :param dataArray: The image
        :type dataArray: xarray DataArray
        :return: The image after removing background
        :rtype: xarray DataArray
        """
        res = dataArray - self.get_offset_from_corner(dataArray, **kwargs)
        res.attrs = copy.copy(dataArray.attrs)
        return res
    
    def crop_image(self, dataSet, center=None, span=None, fringeRemoval=False):
        """Crop the image according to the region of interest (ROI).

        :param dataSet: The images
        :type dataSet: xarray DataArray or DataSet
        :param center: The center of region of insterest (ROI), defaults to None
        :type center: tuple, optional
        :param span: the span of region of insterest (ROI), defaults to None
        :type span: tuple, optional
        :param fringeRemoval: If also crop the reference background images for finges removal function, defaults to False
        :type fringeRemoval: bool, optional
        :return: The croped images
        :rtype: xarray DataArray or DataSet
        """

        if center is None:
            center = self._center
        if span is None:
            span = self._span

        x_start = int(center[0] - span[0] / 2)
        x_end = int(center[0] + span[0] / 2)
        y_end = int(center[1] + span[1] / 2)
        y_start = int(center[1] - span[1] / 2)
        
        dataSet.attrs['x_start'] = x_start
        dataSet.attrs['x_end'] = x_end
        dataSet.attrs['y_end'] = y_end
        dataSet.attrs['y_start'] = y_start
        dataSet.attrs['x_center'] = center[0]
        dataSet.attrs['y_center'] = center[1]
        dataSet.attrs['x_span'] = span[0]
        dataSet.attrs['y_span'] = span[1]
        
        if isinstance(dataSet, type(xr.Dataset())):
            for key in list(dataSet.data_vars):
                dataSet[key].attrs['x_start'] = x_start
                dataSet[key].attrs['x_end'] = x_end
                dataSet[key].attrs['y_end'] = y_end
                dataSet[key].attrs['y_start'] = y_start
                dataSet[key].attrs['x_center'] = center[0]
                dataSet[key].attrs['y_center'] = center[1]
                dataSet[key].attrs['x_span'] = span[0]
                dataSet[key].attrs['y_span'] = span[1]
        
        if fringeRemoval:
            scanAxis = list(get_scanAxis(self.fringeRemovalReferenceImages))
            if not scanAxis[1] is None:
                self._fringeRemoval.referenceImages = self.fringeRemovalReferenceImages.isel(x=slice(x_start, x_end), y=slice(y_start, y_end)).stack(_imgIdx=scanAxis)
            else:
                self._fringeRemoval.referenceImages = self.fringeRemovalReferenceImages.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))

        return dataSet.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
    
    def get_OD(self, imageAtom, imageBackground, imageDrak):
        """Calculate the OD image for absorption imaging.

        :param imageAtom: The image with atoms
        :type imageAtom: numpy array
        :param imageBackground: The image without atoms
        :type imageBackground: numpy array
        :param imageDrak: The image without light
        :type imageDrak: numpy array
        :return: The OD images
        :rtype: numpy array
        """

        numerator = np.atleast_1d(imageBackground - imageDrak)
        denominator = np.atleast_1d(imageAtom - imageDrak)
        
        numerator[numerator == 0] = 1
        denominator[denominator == 0] = 1
        imageOD = np.abs(np.divide(denominator, numerator)) 
        imageOD= -np.log(imageOD)

        if len(imageOD) == 1:
            return imageOD[0]
        else:
            return imageOD
    
    def get_OD_no_dark(self, imageAtom, imageBackground):
        """Calculate the OD image for absorption imaging without dark images.

        :param imageAtom: The image with atoms
        :type imageAtom: numpy array
        :param imageBackground: The image without atoms
        :type imageBackground: numpy array
        :param imageDrak: The image without light
        :type imageDrak: numpy array
        :return: The OD images
        :rtype: numpy array
        """

        numerator = np.atleast_1d(imageBackground)
        denominator = np.atleast_1d(imageAtom)
        
        numerator[numerator == 0] = 1
        denominator[denominator == 0] = 1
        imageOD = np.abs(np.divide(denominator, numerator)) 
        imageOD= -np.log(imageOD)

        if len(imageOD) == 1:
            return imageOD[0]
        else:
            return imageOD

    def get_Ncount(self, dataSet, dim=['x', 'y'], **kwargs):
        """Sum all the value in the image to give the Ncount.

        :param dataSet: The images
        :type dataSet: xarray DataArray or DataSet
        :param dim: The dimensions to take the sumation, defaults to ['x', 'y']
        :type dim: list, optional
        :return: The Ncount
        :rtype: xarray DataArray or DataSet
        """
        return dataSet.sum(dim=['x', 'y'], **kwargs)

    def get_absorption_images(self, dataSet, fringeRemoval=False, dask='allowed', keep_attrs=True, **kwargs):
        """Calculate the OD images for absorption imaging.

        :param dataSet: The data from absorption imaging.
        :type dataSet: xarray DataSet
        :param fringeRemoval: If use fringe removal function, defaults to False
        :type fringeRemoval: bool, optional
        :param dask: over write of the same argument in xarray.apply_ufunc, defaults to 'allowed'
        :type dask: str, optional
        :param keep_attrs: over write of the same argument in xarray.apply_ufunc, defaults to True
        :type keep_attrs: bool, optional
        :return: The data including the OD images.
        :rtype: xarray DataSet
        """

        kwargs.update(
            {
                'dask': dask,
                'keep_attrs': keep_attrs,
            }
        )

        if fringeRemoval:
            
            dataSetAtoms = dataSet[self._image_name['atoms']] - dataSet[self._image_name['dark']]
            
            scanAxis = list(get_scanAxis(dataSet))
            if not scanAxis[1] is None:
                OptimumRef = self._fringeRemoval.fringe_removal(dataSetAtoms.stack(_imgIdx=scanAxis))
            else:
                OptimumRef = self._fringeRemoval.fringe_removal(dataSetAtoms)
            
            dataSet = dataSet.assign(
                {
                    self._image_name['optimumBackground']: OptimumRef
                }
            )
            
            dataSet = dataSet.assign(
                {
                    self._image_name['OD']: xr.apply_ufunc(self.get_OD_no_dark, dataSetAtoms, dataSet[self._image_name['optimumBackground']], **kwargs)
                }
            )
        else:
            dataSet = dataSet.assign(
                {
                    self._image_name['OD']: xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)
                }
            )
        
        # dataSet[self._image_name['OD']].attrs.update(dataSet.attrs)

        return dataSet

    def remove_background(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):
        """Remove the background for the given image data

        :param dataSet: The data from absorption imaging.
        :type dataSet: xarray DataSet
        :param dask: over write of the same argument in xarray.apply_ufunc, defaults to 'allowed'
        :type dask: str, optional
        :param keep_attrs: over write of the same argument in xarray.apply_ufunc, defaults to True
        :type keep_attrs: bool, optional
        """
        
        kwargs.update(
            {
                'dask': dask,
                'keep_attrs': keep_attrs,
            }
        )

        xr.apply_ufunc(self.get_OD, dataSet[self._image_name['atoms']], dataSet[self._image_name['background']], dataSet[self._image_name['dark']], **kwargs)
    
    @property
    def fringeRemovalCenter(self):
        """The getter of the center of region of insterest (ROI)

        :return: The center of region of insterest (ROI)
        :rtype: tuple
        """
        return self._fringeRemoval.center
    
    @fringeRemovalCenter.setter
    def fringeRemovalCenter(self, value):
        """The setter of the center of region of insterest (ROI)

        :param value: The center of region of insterest (ROI)
        :type value: tuple
        """
        self._fringeRemoval.center = value

    @property
    def fringeRemovalSpan(self):
        """The getter of the span of region of insterest (ROI) 

        :return: The span of region of insterest (ROI) 
        :rtype: tuple
        """
        return self._fringeRemoval.span

    @fringeRemovalSpan.setter
    def fringeRemovalSpan(self, value):
        """The setter of the span of region of insterest (ROI) 

        :param value: The span of region of insterest (ROI) 
        :type value: tuple
        """
        self._fringeRemoval.span = value
        
    def load_fringe_removal_background_from_hdf5(self, img_dir, SequenceName, data, shotNum, group, crop=False, load=False, **kwargs):
        folderPath = img_dir + SequenceName + "/" + data
        filePath = folderPath + "/" + shotNum + "/*.h5"
        
        dataSet = read_hdf5_file(filePath, group, **kwargs)
        scanAxis = dataSet.scanAxis
        
        dataSet = dataSet[self._image_name['background']] - dataSet[self._image_name['dark']]
        dataSet.attrs['scanAxis'] = scanAxis
        
        if crop:
            dataSet = self.crop_image(dataSet)
        
        if load:
            self.fringeRemovalReferenceImages = dataSet.load()
        else:
            self.fringeRemovalReferenceImages = dataSet
    
    def load_fringe_removal_background_from_database():
        pass