161 lines
5.2 KiB
Python
161 lines
5.2 KiB
Python
import numpy as np
|
|
import xarray as xr
|
|
import copy
|
|
|
|
|
|
class ImageAnalyser():
|
|
|
|
def __init__(self) -> None:
|
|
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):
|
|
return self._image_name
|
|
|
|
@image_name.setter
|
|
def image_name(self, value):
|
|
self._image_name.update(value)
|
|
|
|
@property
|
|
def center(self):
|
|
return self._center
|
|
|
|
@center.setter
|
|
def center(self, value):
|
|
self._center = value
|
|
|
|
@property
|
|
def span(self):
|
|
return self._span
|
|
|
|
@span.setter
|
|
def span(self, value):
|
|
self._span = value
|
|
|
|
@property
|
|
def fraction(self):
|
|
return self._fraction
|
|
|
|
@fraction.setter
|
|
def fraction(self, value):
|
|
self._fraction = value
|
|
|
|
def get_offset_from_corner(self, dataArray, x_fraction=None, y_fraction=None, fraction=None, xAxisName='x', yAxisName='y'):
|
|
|
|
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):
|
|
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):
|
|
|
|
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]
|
|
|
|
return dataSet.isel(x=slice(x_start, x_end), y=slice(y_start, y_end))
|
|
|
|
def get_OD(self, imageAtom, imageBackground, imageDrak):
|
|
|
|
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):
|
|
return dataSet.sum(dim=['x', 'y'], **kwargs)
|
|
|
|
def get_absorption_images(self, dataSet, dask='allowed', keep_attrs=True, **kwargs):
|
|
|
|
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):
|
|
|
|
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)
|
|
|
|
|
|
|