src/libertem/masks.py
from typing import Union, Callable
from collections.abc import Iterable
import numpy as np
import scipy.sparse as sp
import sparse
from libertem.utils import make_polar
# Import here for backwards compatibility, refs #1031
from libertem.common.sparse import to_dense, to_sparse, is_sparse # NOQA: F401
MaskArrayType = Union[np.ndarray, sp.coo_matrix, sp.dok_matrix]
MaskFactoriesType = Union[Callable[[], MaskArrayType], Iterable[Callable[[], MaskArrayType]]]
def _make_circular_mask(centerX, centerY, imageSizeX, imageSizeY, radius, antialiased=False):
"""
Make a circular mask in a bool array for masking a region in an image.
Parameters
----------
centerX, centerY : float
Centre point of the mask.
imageSizeX, imageSizeY : int
Size of the image to be masked.
radius : float
Radius of the mask.
Returns
-------
Boolean Numpy 2D Array
Array with the shape (imageSizeX, imageSizeY) with the mask.
Examples
--------
>>> image = np.ones((9, 9))
>>> mask = _make_circular_mask(4, 4, 9, 9, 2)
>>> image_masked = image*mask
>>> import matplotlib.pyplot as plt
>>> cax = plt.imshow(image_masked)
"""
if antialiased:
mask = radial_bins(
centerX, centerY, imageSizeX, imageSizeY, radius, n_bins=1, use_sparse=False
)[0]
else:
x, y = np.ogrid[-centerY:imageSizeY-centerY, -centerX:imageSizeX-centerX]
mask = x*x + y*y <= radius*radius
return mask
def sparse_template_multi_stack(mask_index, offsetX, offsetY, template, imageSizeX, imageSizeY):
'''
Stamp the template in a multi-mask 3D stack at the positions indicated by
mask_index, offsetY, offsetX. The function clips the bounding box as necessary.
'''
num_templates = len(mask_index)
fy, fx = template.shape
area = fy * fx
total_index_size = num_templates * area
y, x = np.mgrid[0:fy, 0:fx]
data = np.zeros(total_index_size, dtype=template.dtype)
coord_mask = np.zeros(total_index_size, dtype=int)
coord_y = np.zeros(total_index_size, dtype=int)
coord_x = np.zeros(total_index_size, dtype=int)
for i in range(len(mask_index)):
start = i * area
stop = (i + 1) * area
data[start:stop] = template.flatten()
coord_mask[start:stop] = mask_index[i]
coord_y[start:stop] = y.flatten() + offsetY[i]
coord_x[start:stop] = x.flatten() + offsetX[i]
selector = (coord_y >= 0) * (coord_y < imageSizeY) * (coord_x >= 0) * (coord_x < imageSizeX)
return sparse.COO(
data=data[selector],
coords=(coord_mask[selector], coord_y[selector], coord_x[selector]),
shape=(int(max(mask_index) + 1), imageSizeY, imageSizeX)
)
def sparse_circular_multi_stack(mask_index, centerX, centerY, imageSizeX, imageSizeY, radius):
# we make sure it is odd
bbox = int(2*np.ceil(radius) + 1)
bbox_center = int((bbox - 1) // 2)
template = circular(
centerX=bbox_center,
centerY=bbox_center,
imageSizeX=bbox,
imageSizeY=bbox,
radius=radius)
return sparse_template_multi_stack(
mask_index=mask_index,
offsetX=np.array(centerX, dtype=int) - bbox_center,
offsetY=np.array(centerY, dtype=int) - bbox_center,
template=template,
imageSizeX=imageSizeX,
imageSizeY=imageSizeY,
)
def circular(centerX, centerY, imageSizeX, imageSizeY, radius, antialiased=False):
"""
Make a circular mask as a 2D array
Parameters
----------
centreX, centreY : float
Centre point of the mask.
imageSizeX, imageSizeY : int
Size of the image to be masked.
radius : float
Radius of the mask.
Returns
-------
Numpy 2D Array
Array with the shape (imageSizeX, imageSizeY) with the mask.
"""
mask = _make_circular_mask(centerX, centerY, imageSizeX, imageSizeY, radius, antialiased)
return mask
def ring(centerX, centerY, imageSizeX, imageSizeY, radius, radius_inner, antialiased=False):
"""
Make a ring mask as a double array.
Parameters
----------
centreX, centreY : float
Centre point of the mask.
imageSizeX, imageSizeY : int
Size of the image to be masked.
radius : float
Outer radius of the ring.
radius_inner : float
Inner radius of the ring.
Returns
-------
Numpy 2D Array
Array with the shape (imageSizeX, imageSizeY) with the mask.
"""
if antialiased:
mask = radial_bins(
centerX, centerY, imageSizeX, imageSizeY,
radius=radius, radius_inner=radius_inner, n_bins=1, use_sparse=False
)[0]
else:
outer = _make_circular_mask(centerX, centerY, imageSizeX, imageSizeY, radius)
inner = _make_circular_mask(centerX, centerY, imageSizeX, imageSizeY, radius_inner)
mask = outer & ~inner
return mask
def radial_gradient(centerX, centerY, imageSizeX, imageSizeY, radius, antialiased=False):
'''
Generate a linear radial gradient from 0 to 1 within radius
'''
x, y = np.ogrid[-centerY:imageSizeY-centerY, -centerX:imageSizeX-centerX]
if antialiased:
r = np.sqrt(x**2 + y**2)
mask = radial_gradient_background_subtraction(
r=r, r0=radius, r_outer=0
)
else:
mask = (x*x + y*y <= radius*radius) * (np.sqrt(x*x + y*y) / radius)
return mask
def radial_gradient_background_subtraction(r, r0, r_outer, delta=1):
'''
Generate a template with a linear radial gradient from 0 to 1 inside r0,
linear transition region for antialiasing between [r0 - delta/2, r0 + delta/2[,
and a negative ring with value -1 in [r0 + delta/2, r_outer].
The function accepts the radius for each pixel as a parameter so that a distorted version can
be generated with the stretchY and angle parameters of :meth:`~libertem.masks.polar_map`.
Parameters
----------
r : numpy.ndarray
Map of radius for each pixel, typically 2D. This allows to work in distorted coordinate
systems by assigning arbitrary radius values to each pixel.
:meth:`~libertem.masks.polar_map` can generate elliptical maps as an example.
r0 : float
Inner radius to fill with a linear gradient in units of r
r_outer : float
Outer radius of ring from r0 to fill with -1 in units of r
delta : float, optional
Width of transition region between inner and outer in units of r
with linear gradient for antialiasing or smoothening. Defaults to 1.
Returns
-------
numpy.ndarray
NumPy numpy.ndarray with the same shape and type of r with mask values assigned as
described in the description.
'''
result = np.zeros_like(r)
within = r < r0 - delta/2
result[within] = r[within] / r0
transition = (r >= r0 - delta/2) * (r < r0 + delta/2)
result[transition] = (r0 - r[transition]) / (delta/2)
without = (r >= r0 + delta/2) * (r <= r_outer)
result[without] = -1
return result
def polar_map(centerX, centerY, imageSizeX, imageSizeY, stretchY=1., angle=0.):
'''
Return a map of radius and angle.
The optional parameters stretchY and angle allow to stretch and rotate the coordinate system
into an elliptical form. This is useful to generate modified input data for functions that
generate a template as a function of radius and angle.
Parameters
----------
centerX,centerY : float
Center of the coordinate system in pixel coordinates
imageSizeX,imageSizeY : int
Size of the map to generate in px
stretchY,angle : float, optional
Stretch the radius elliptically by amount :code:`stretchY` in direction
:code:`angle` in radians. :code:`angle = 0` means in Y direction.
Returns
-------
Tuple[numpy.ndarray, numpy.ndarray]
Map of radius and angle of shape :code:`(imageSizeY, imageSizeX)`
'''
y, x = np.mgrid[0:imageSizeY, 0:imageSizeX]
dy = y - centerY
dx = x - centerX
if stretchY != 1.0 or angle != 0.:
(dy, dx) = (
(dy*np.cos(angle) - dx*np.sin(angle)) / stretchY,
dx*np.cos(angle) + dy*np.sin(angle),
)
dy = dy.flatten()
dx = dx.flatten()
cartesians = np.stack((dy, dx)).T
polars = make_polar(cartesians)
return (
polars[:, 0].reshape((imageSizeY, imageSizeX)),
polars[:, 1].reshape((imageSizeY, imageSizeX))
)
def balance(template):
'''
Accept a template with both positive and negative values and scale the negative
part in such a way that the sum is zero.
This is useful to generate masks that return zero when applied to a
uniform background or linear gradient.
'''
result = template.copy()
above = template > 0
below = template < 0
result[below] *= template[above].sum() / template[below].sum() * -1
return result
def bounding_radius(centerX, centerY, imageSizeX, imageSizeY):
'''
Calculate a radius around centerX, centerY that covers the whole frame
'''
dy = max(centerY, imageSizeY - centerY)
dx = max(centerX, imageSizeX - centerX)
return int(np.ceil(np.sqrt(dy**2 + dx**2))) + 1
def radial_bins(centerX, centerY, imageSizeX, imageSizeY,
radius=None, radius_inner=0, n_bins=None, normalize=False, use_sparse=None, dtype=None):
'''
Generate antialiased rings
'''
if radius is None:
radius = bounding_radius(centerX, centerY, imageSizeX, imageSizeY)
if n_bins is None:
n_bins = int(np.round(radius - radius_inner))
r, phi = polar_map(centerX, centerY, imageSizeX, imageSizeY)
r = r.flatten()
width = (radius - radius_inner) / n_bins
bin_area = np.pi * (radius**2 - (radius - width)**2)
if use_sparse is None:
use_sparse = bin_area / (imageSizeX * imageSizeY) < 0.1
if use_sparse:
jjs = np.arange(len(r), dtype=np.int64)
slices = []
for r0 in np.linspace(radius_inner, radius - width, n_bins) + width/2:
diff = np.abs(r - r0)
# The "0.5" ensures that the bins overlap and sum up to exactly 1
vals = np.maximum(0, np.minimum(1, width/2 + 0.5 - diff))
if use_sparse:
select = vals != 0
vals = vals[select]
if normalize: # Make sure each bin has a sum of 1
s = vals.sum()
if not np.isclose(s, 0):
vals /= s
slices.append(sparse.COO(shape=len(r), data=vals.astype(dtype), coords=(jjs[select],)))
else:
if normalize: # Make sure each bin has a sum of 1
s = vals.sum()
if not np.isclose(s, 0):
vals /= s
slices.append(vals.reshape((imageSizeY, imageSizeX)).astype(dtype))
# Patch a singularity at the center
if radius_inner < 0.5:
yy = int(np.round(centerY))
xx = int(np.round(centerX))
if yy >= 0 and yy < imageSizeY and xx >= 0 and xx < imageSizeX:
if use_sparse:
index = yy * imageSizeX + xx
diff = 1 - slices[0][index] - radius_inner
patch = sparse.COO(shape=len(r), data=[diff], coords=[index])
slices[0] += patch
else:
slices[0][yy, xx] = 1 - radius_inner
if use_sparse:
return sparse.stack(slices).reshape((-1, imageSizeY, imageSizeX))
else:
return np.stack(slices)
def background_subtraction(centerX, centerY, imageSizeX, imageSizeY, radius, radius_inner,
antialiased=False):
mask_1 = circular(
centerX, centerY, imageSizeX, imageSizeY, radius_inner, antialiased=antialiased
)
sum_1 = np.sum(mask_1)
mask_2 = ring(
centerX, centerY, imageSizeX, imageSizeY, radius, radius_inner, antialiased=antialiased
)
sum_2 = np.sum(mask_2)
mask = mask_1 - mask_2*sum_1/sum_2
return mask
def rectangular(X, Y, Width, Height, imageSizeX, imageSizeY):
"""
Make a rectangular mask as a 2D array of bool.
Parameters
----------
X, Y : Corner coordinates
Centre point of the mask.
imageSizeX, imageSizeY : int
Size of the image to be masked.
Width, Height : Width and Height of the rectangle
Returns
-------
Numpy 2D Array
Array with the shape (imageSizeX, imageSizeY) with the mask.
"""
bool_mask = np.zeros([imageSizeY, imageSizeX], dtype="bool")
if Height*Width > 0:
ymin = min(Y, Y+Height)
xmin = min(X, X+Width)
ymax = max(Y, Y+Height)
xmax = max(X, X+Width)
elif Height > 0 and Width < 0:
ymin = Y
xmin = X+Width
ymax = Y+Height
xmax = X
elif Height < 0 and Width > 0:
ymin = Y+Height
xmin = X
ymax = Y
xmax = X+Width
else:
ymin = 0
xmin = 0
ymax = -1
xmax = -1
ymin = int(ymin)
xmin = int(xmin)
ymax = int(ymax)
xmax = int(xmax)
bool_mask[max(0, ymin):min(ymax+1, imageSizeY), max(0, xmin):min(xmax+1, imageSizeX)] = 1
return bool_mask
# TODO: dtype parameter? consistency with ring/circular above
def gradient_x(imageSizeX, imageSizeY, dtype=np.float32):
return np.tile(
np.ogrid[slice(0, imageSizeX)].astype(dtype), imageSizeY
).reshape(imageSizeY, imageSizeX)
def gradient_y(imageSizeX, imageSizeY, dtype=np.float32):
return gradient_x(imageSizeY, imageSizeX, dtype).transpose()