From 87d6ff49c98c0d70c3636304a0061f0bb9d7c357 Mon Sep 17 00:00:00 2001 From: Gao Date: Thu, 28 Sep 2023 16:35:30 +0200 Subject: [PATCH] optimize the fringes removal function --- Analyser/FringeRemoval.py | 10 ++- Analyser/ImagingAnalyser.py | 141 +++++++++++++++++++++++++++++++++--- 2 files changed, 141 insertions(+), 10 deletions(-) diff --git a/Analyser/FringeRemoval.py b/Analyser/FringeRemoval.py index b675569..2da44b9 100644 --- a/Analyser/FringeRemoval.py +++ b/Analyser/FringeRemoval.py @@ -195,6 +195,10 @@ class FringeRemoval(): @referenceImages.setter def referenceImages(self, value): + if value is None: + self._referenceImages = None + return + if self.reshape: value = self.reshape_data(value) elif isinstance(value, type(xr.DataArray())): @@ -257,6 +261,10 @@ class FringeRemoval(): raise DataSizeException self._mask = value + + self._center = None + self._span = None + self.k = np.where(self._mask.flatten() == 1)[0] def _auto_mask(self): @@ -267,7 +275,7 @@ class FringeRemoval(): y_end = int(self._center[1] + self._span[1] / 2) y_start = int(self._center[1] - self._span[1] / 2) - mask[y_start:y_end, x_start:x_end] = 0 + mask[y_start:y_end, x_start:x_end] = 0 return mask diff --git a/Analyser/ImagingAnalyser.py b/Analyser/ImagingAnalyser.py index 488e68c..7359bc1 100644 --- a/Analyser/ImagingAnalyser.py +++ b/Analyser/ImagingAnalyser.py @@ -2,6 +2,10 @@ 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 @@ -15,11 +19,15 @@ class ImageAnalyser(): '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 @@ -150,7 +158,7 @@ class ImageAnalyser(): res.attrs = copy.copy(dataArray.attrs) return res - def crop_image(self, dataSet, center=None, span=None): + 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 @@ -192,6 +200,13 @@ class ImageAnalyser(): 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)) @@ -220,6 +235,32 @@ class ImageAnalyser(): 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. @@ -233,7 +274,7 @@ class ImageAnalyser(): """ return dataSet.sum(dim=['x', 'y'], **kwargs) - def get_absorption_images(self, dataSet, dask='allowed', keep_attrs=True, **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. @@ -253,11 +294,31 @@ class ImageAnalyser(): } ) - 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) - } - ) + if fringeRemoval: + + scanAxis = list(get_scanAxis(dataSet)) + if not scanAxis[1] is None: + OptimumRef = self._fringeRemoval.fringe_removal(dataSet[self._image_name['atoms']].stack(_imgIdx=scanAxis)) + else: + OptimumRef = self._fringeRemoval.fringe_removal(dataSet[self._image_name['atoms']]) + + dataSet = dataSet.assign( + { + self._image_name['optimumBackground']: OptimumRef + } + ) + + dataSet = dataSet.assign( + { + self._image_name['OD']: xr.apply_ufunc(self.get_OD_no_dark, dataSet[self._image_name['atoms']], 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) @@ -282,6 +343,68 @@ class ImageAnalyser(): ) 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) - \ No newline at end of file + :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 + + # @property + # def fringeRemovalReferenceImages(self): + # return self._fringeRemoval.referenceImages + + # @fringeRemovalReferenceImages.setter + # def fringeRemovalReferenceImages(self, value): + # self._fringeRemoval.referenceImages = 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.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 \ No newline at end of file