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.
296 lines
10 KiB
296 lines
10 KiB
import numpy as np
|
|
import xarray as xr
|
|
import copy
|
|
|
|
|
|
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',
|
|
}
|
|
self._center = None
|
|
self._span = None
|
|
self._fraction = 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):
|
|
"""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
|
|
:return: The croped images
|
|
:rtype: xarray DataArray or DataSet
|
|
"""
|
|
|
|
if center is None:
|
|
center = self._center
|
|
if span is None:
|
|
span = self._span
|
|
|
|
if not x in
|
|
|
|
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]
|
|
|
|
res = dataSet.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
|
|
res = res.assign_coords(
|
|
{
|
|
'x': np.linspace(x_start, x_end - 1, span[0]),
|
|
'y': np.linspace(y_start, y_end - 1, span[1]),
|
|
}
|
|
)
|
|
return res
|
|
|
|
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_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, 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 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,
|
|
}
|
|
)
|
|
|
|
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)
|
|
|
|
|
|
|