def get_mask_from_roi(signal, roi, axes=None, gaussian=True):
if axes is None and signal in roi.signal_map:
axes = roi.signal_map[signal][1]
else:
axes = roi._parse_axes(axes, signal.axes_manager)
# Needs to add support for other type of ROI
if hasattr(roi, 'cx'):
# CircleROI
radius = roi.r
cx = roi.cx
cy = roi.cy
r = np.linalg.norm([cx, cy]) * 0.8
# The factor of 3 come from an estimate of how far the tail of the
# Gaussian goes; to avoid getting the zero-frequency component in
# the mask, we clip its radius_slice value
radius_slice = np.clip(radius * 3, a_min=radius, a_max=r)
ranges = [[cx - radius_slice, cx + radius_slice],
[cy - radius_slice, cy + radius_slice]]
else:
ranges = roi._get_ranges()
if hasattr(roi, 'cx'):
# The 'else' part is missing
slices = roi._make_slices(axes, axes, ranges=ranges)
if not gaussian:
# in case of Bragg Filtering
radius_slice = radius
# Calculate a disk mask
sig_axes = signal.axes_manager.signal_axes
ir = [slices[sig_axes.index(axes[0])], slices[sig_axes.index(axes[1])]]
vx = axes[0].axis[ir[0]] - cx
vy = axes[1].axis[ir[1]] - cy
gx, gy = np.meshgrid(vx, vy)
gr = gx**2 + gy**2
disk_mask = gr > radius_slice**2
if gaussian:
import hyperspy.api as hs
mask = hs.signals.Signal2D(np.zeros(signal.data.shape))
x = np.linspace(ranges[0][0], ranges[0][1], disk_mask.shape[1])
y = np.linspace(ranges[1][0], ranges[1][1], disk_mask.shape[0])
xx, yy = np.meshgrid(x, y)
gaussian2d = hs.model.components2D.Gaussian2D(sigma_x=radius,
sigma_y=radius,
centre_x=cx,
centre_y=cy,
A=2 * np.pi *
radius**2)
mask_circle = gaussian2d.function(xx, yy) * ~disk_mask
else:
mask = BaseSignal(np.full(signal.data.shape, True, dtype=bool))
mask.axes_manager.set_signal_dimension(
signal.axes_manager.signal_dimension)
mask_circle = disk_mask
mask.isig[slices] = mask_circle
# If signal.data is cupy array, transfer the array to the GPU
xp = get_array_module(signal.data)
mask.data = xp.asarray(mask.data)
return mask
