You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
425 lines
16 KiB
425 lines
16 KiB
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, date, shotNum, group, crop=False, load=False, **kwargs):
|
|
"""Load the reference background images from hdf5 files of one single shot.
|
|
|
|
:param img_dir: The path of the folder storing data.
|
|
:type img_dir: str
|
|
:param SequenceName: The name of the sequence
|
|
:type SequenceName: str
|
|
:param date: The date when the shot was taken in 'YYYY/MM/DD'.
|
|
:type date: str
|
|
:param shotNum: The number of the shot
|
|
:type shotNum: str
|
|
:param group: The name of the group storing the imgaes
|
|
:type group: str
|
|
:param crop: If crop the data, defaults to False
|
|
:type crop: bool, optional
|
|
:param load: If load the data into RAM, defaults to False
|
|
:type load: bool, optional
|
|
"""
|
|
folderPath = img_dir + SequenceName + "/" + date
|
|
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
|