from __future__ import annotations
from abc import ABC, abstractmethod
import numpy as np
from skimage import draw, filters
from skimage.draw import polygon
from ..array_utils import geometric_center_idx
def clip_add(
image1: np.ndarray, image2: np.ndarray, dtype: type[np.dtype] = np.uint16
) -> np.ndarray:
"""Clip the image to the dtype extrema. Otherwise, the bits will flip."""
# convert to float first so we don't flip bits initially
combined_img = image1.astype(float) + image2.astype(float)
return np.clip(combined_img, np.iinfo(dtype).min, np.iinfo(dtype).max).astype(dtype)
def clip_multiply(
image1: np.ndarray, image2: np.ndarray, dtype: type[np.dtype] = np.uint16
) -> np.ndarray:
"""Clip the image to the dtype extrema. Otherwise, the bits will flip."""
# convert to float first so we don't flip bits initially
combined_img = image1.astype(float) * image2.astype(float)
return np.clip(combined_img, np.iinfo(dtype).min, np.iinfo(dtype).max).astype(dtype)
def even_round(num: float) -> int:
"""Return an even number"""
num = int(round(num))
return num + num % 2
def gaussian2d(
mx: float,
my: float,
height: float,
center_x: float,
center_y: float,
width_x: float,
width_y: float,
constant: float = 0,
) -> np.ndarray:
"""Returns a gaussian function with the given parameters"""
width_x = float(width_x)
width_y = float(width_y)
return (
height
* np.exp(
-(((center_x - mx) / width_x) ** 2 + ((center_y - my) / width_y) ** 2) / 2
)
+ constant
)
class Layer(ABC):
"""Abstract base for layers"""
@abstractmethod
def apply(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> np.ndarray:
"""Apply the layer. Takes a 2D array and pixel size value in and returns a modified array."""
pass
[docs]
class PerfectConeLayer(Layer):
"""A cone without flattening filter effects"""
def __init__(
self,
cone_size_mm: float = 10,
cax_offset_mm: (float, float) = (0, 0),
alpha: float = 1.0,
rotation: float = 0,
):
"""
Parameters
----------
cone_size_mm
Cone size in mm at the iso plane
cax_offset_mm
The offset in mm. (down, right)
alpha
The intensity of the layer. 1 is full saturation/radiation. 0 is none.
rotation: float
The amount of rotation in degrees. When there is an offset, this acts like a couch kick.
"""
self.cone_size_mm = cone_size_mm
self.cax_offset_mm = cax_offset_mm
self.alpha = alpha
self.rotation = rotation
[docs]
def apply(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> np.ndarray:
image, _, _ = self._create_perfect_field(image, pixel_size, mag_factor)
return image
def _create_perfect_field(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> (np.ndarray, ...):
cone_size_pix = ((self.cone_size_mm / 2) / pixel_size) * mag_factor**2
# we rotate the point around the center of the image
offset_pix_y, offset_pix_x = rotate_point(
x=self.cax_offset_mm[0] * mag_factor / pixel_size,
y=self.cax_offset_mm[1] * mag_factor / pixel_size,
angle=self.rotation,
)
# convert to pixels and shift to center
cax_offset_pix = (
offset_pix_y + (image.shape[0] / 2 - 0.5),
offset_pix_x + (image.shape[1] / 2 - 0.5),
)
rr, cc = draw.disk(cax_offset_pix, cone_size_pix, shape=image.shape)
rr = np.round(rr).astype(int)
cc = np.round(cc).astype(int)
temp_array = np.zeros(image.shape)
temp_array[rr, cc] = int(np.iinfo(image.dtype).max * self.alpha)
image = clip_add(image, temp_array)
return image, rr, cc
[docs]
class FilterFreeConeLayer(PerfectConeLayer):
"""A cone with flattening filter effects."""
def __init__(
self,
cone_size_mm: float = 10,
cax_offset_mm: (float, float) = (0, 0),
alpha: float = 1.0,
filter_magnitude: float = 0.4,
filter_sigma_mm: float = 80,
):
"""
Parameters
----------
cone_size_mm
Cone size in mm at the iso plane
cax_offset_mm
The offset in mm. (out, right)
alpha
The intensity of the layer. 1 is full saturation/radiation. 0 is none.
filter_magnitude
The magnitude of the CAX peak. Larger values result in "pointier" fields.
filter_sigma_mm
Proportional to the width of the CAX peak. Larger values produce wider curves.
"""
super().__init__(cone_size_mm, cax_offset_mm, alpha)
self.filter_magnitude = filter_magnitude
self.filter_sigma_mm = filter_sigma_mm
[docs]
def apply(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> np.ndarray:
image, rr, cc = self._create_perfect_field(image, pixel_size, mag_factor)
# add filter effect
center_x = geometric_center_idx(image[:, 0])
center_y = geometric_center_idx(image[0, :])
n = gaussian2d(
rr,
cc,
self.filter_magnitude * np.iinfo(image.dtype).max,
center_x,
center_y,
self.filter_sigma_mm / pixel_size,
self.filter_sigma_mm / pixel_size,
constant=-self.filter_magnitude * np.iinfo(image.dtype).max,
)
image[rr, cc] += n.astype(image.dtype)
return image
[docs]
class PerfectFieldLayer(Layer):
"""A square field without flattening filter effects"""
def __init__(
self,
field_size_mm: (float, float) = (10, 10),
cax_offset_mm: (float, float) = (0, 0),
alpha: float = 1.0,
rotation: float = 0,
):
"""
Parameters
----------
field_size_mm
Field size in mm at the iso plane as (height, width)
cax_offset_mm
The offset in mm. (down, right)
alpha
The intensity of the layer. 1 is full saturation/radiation. 0 is none.
rotation: float
The amount of rotation in degrees. This acts like a collimator rotation.
"""
self.field_size_mm = field_size_mm
self.cax_offset_mm = cax_offset_mm
self.alpha = alpha
self.rotation = rotation
def _create_perfect_field(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> (np.ndarray, ...):
field_size_pix = [
even_round(f * mag_factor**2 / pixel_size) for f in self.field_size_mm
]
cax_offset_pix_mag = [v * mag_factor / pixel_size for v in self.cax_offset_mm]
field_center = [
offset + (shape / 2) - 0.5
for offset, shape in zip(cax_offset_pix_mag, image.shape)
]
rr, cc = draw_rotated_rectangle(
image.shape, center=field_center, extent=field_size_pix, angle=self.rotation
)
rr = np.round(rr).astype(int)
cc = np.round(cc).astype(int)
temp_array = np.zeros(image.shape)
temp_array[rr, cc] = int(np.iinfo(image.dtype).max * self.alpha)
image = clip_add(image, temp_array)
return image, rr, cc
[docs]
def apply(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> np.array:
image, _, _ = self._create_perfect_field(image, pixel_size, mag_factor)
return image
[docs]
class FilteredFieldLayer(PerfectFieldLayer):
"""A square field with flattening filter effects"""
def __init__(
self,
field_size_mm: (float, float) = (10, 10),
cax_offset_mm: (float, float) = (0, 0),
alpha: float = 1.0,
gaussian_height: float = 0.03,
gaussian_sigma_mm: float = 32,
rotation: float = 0,
):
"""
Parameters
----------
field_size_mm
Field size in mm at the iso plane (height, width)
cax_offset_mm
The offset in mm. (out, right)
alpha
The intensity of the layer. 1 is full saturation/radiation. 0 is none.
gaussian_height
The intensity of the "horns", or more accurately, the CAX dip. Proportional to the max value allowed for the data type.
Increase to make the horns more prominent.
gaussian_sigma_mm
The width of the "horns". A.k.a. the CAX dip width. Increase to create a wider
horn effect.
rotation: float
The amount of rotation in degrees. This acts like a collimator rotation.
"""
super().__init__(
field_size_mm=field_size_mm,
cax_offset_mm=cax_offset_mm,
alpha=alpha,
rotation=rotation,
)
self.gaussian_height = gaussian_height
self.gaussian_sigma_mm = gaussian_sigma_mm
[docs]
def apply(self, image: np.array, pixel_size: float, mag_factor: float) -> np.array:
image, rr, cc = self._create_perfect_field(image, pixel_size, mag_factor)
# add filter effect
height = -self.gaussian_height * np.iinfo(image.dtype).max
width = self.gaussian_sigma_mm / pixel_size
center_x = geometric_center_idx(image[:, 0])
center_y = geometric_center_idx(image[0, :])
horns = gaussian2d(
rr,
cc,
height=height,
center_x=center_x,
center_y=center_y,
width_x=width,
width_y=width,
)
# the horns are negative, so we don't have
# to worry about clipping here
image[rr, cc] += horns.astype(image.dtype)
return image
[docs]
class FilterFreeFieldLayer(FilteredFieldLayer):
"""A square field with flattening filter free (FFF) effects"""
def __init__(
self,
field_size_mm: (float, float) = (10, 10),
cax_offset_mm: (float, float) = (0, 0),
alpha: float = 1.0,
gaussian_height: float = 0.4,
gaussian_sigma_mm: float = 80,
rotation: float = 0,
):
"""
Parameters
----------
field_size_mm
Field size in mm at the iso plane (height, width).
cax_offset_mm
The offset in mm. (out, right)
alpha
The intensity of the layer. 1 is full saturation/radiation. 0 is none.
gaussian_height
The magnitude of the CAX peak. Larger values result in "pointier" fields.
gaussian_sigma_mm
Proportional to the width of the CAX peak. Larger values produce wider curves.
rotation: float
The amount of rotation in degrees. This acts like a collimator rotation.
"""
super().__init__(
field_size_mm,
cax_offset_mm,
alpha,
gaussian_height,
gaussian_sigma_mm,
rotation=rotation,
)
[docs]
def apply(self, image: np.array, pixel_size: float, mag_factor: float) -> np.array:
image, rr, cc = self._create_perfect_field(image, pixel_size, mag_factor)
# add filter effect
center_x = geometric_center_idx(image[:, 0])
center_y = geometric_center_idx(image[0, :])
n = gaussian2d(
rr,
cc,
self.gaussian_height * np.iinfo(image.dtype).max,
center_x,
center_y,
self.gaussian_sigma_mm / pixel_size,
self.gaussian_sigma_mm / pixel_size,
constant=-self.gaussian_height * np.iinfo(image.dtype).max,
)
# the horns are negative, so we don't have
# to worry about clipping here
image[rr, cc] += n.astype(image.dtype)
return image
[docs]
class PerfectBBLayer(PerfectConeLayer):
"""A BB-like layer. Like a cone, but with lower alpha (i.e. higher opacity)"""
def __init__(
self,
bb_size_mm: float = 5,
cax_offset_mm: (float, float) = (0, 0),
alpha: float = -0.5,
rotation: float = 0,
):
super().__init__(
cone_size_mm=bb_size_mm,
cax_offset_mm=cax_offset_mm,
alpha=alpha,
rotation=rotation,
)
[docs]
class GaussianFilterLayer(Layer):
"""A Gaussian filter. Simulates the effects of scatter on the field"""
def __init__(self, sigma_mm: float = 2):
self.sigma_mm = sigma_mm
[docs]
def apply(self, image: np.array, pixel_size: float, mag_factor: float) -> np.array:
sigma_pix = self.sigma_mm / pixel_size
return filters.gaussian(image, sigma_pix, preserve_range=True).astype(
image.dtype
)
[docs]
class RandomNoiseLayer(Layer):
"""A salt and pepper noise, simulating dark current"""
def __init__(self, mean: float = 0.0, sigma: float = 0.001):
self.mean = mean
self.sigma = sigma
[docs]
def apply(self, image: np.array, pixel_size: float, mag_factor: float) -> np.array:
normalized_sigma = self.sigma * np.iinfo(image.dtype).max
rng = np.random.default_rng()
noise = rng.normal(self.mean, normalized_sigma, size=image.shape)
return clip_add(image, noise, dtype=image.dtype)
[docs]
class ConstantLayer(Layer):
"""A constant layer. Can be used to simulate scatter or background."""
def __init__(self, constant: float):
self.constant = constant
[docs]
def apply(self, image: np.array, pixel_size: float, mag_factor: float) -> np.array:
constant_img = np.full(image.shape, fill_value=self.constant)
return clip_add(image, constant_img, dtype=image.dtype)
[docs]
class SlopeLayer(Layer):
"""Adds a slope in both directions of the image. Usually used for simulating asymmetry or a-flatness.
Parameters
----------
slope_x : float
The slope in the x-direction (left/right). If positive, will increase the right side.
The value is multiplicative to the current state of the image. E.g. a value of 0.1 will increase the right side by 10% and 0% on the left.
slope_y : float
The slope in the y-direction (up/down). If positive, will increase the bottom side.
"""
def __init__(self, slope_x: float, slope_y: float):
self.slope_x = slope_x
self.slope_y = slope_y
[docs]
def apply(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> np.ndarray:
nrows, ncols = image.shape
# Create the row and column scaling factors
y_scaling = (1 + self.slope_y * np.arange(nrows) / nrows).reshape(-1, 1)
x_scaling = (1 + self.slope_x * np.arange(ncols) / ncols).reshape(1, -1)
y_scaled = clip_multiply(image, y_scaling)
xy_scaled = clip_multiply(y_scaled, x_scaling)
return xy_scaled
[docs]
class ArrayLayer(Layer):
"""Add an already-existing array as a layer. Useful if the array is already
constructed to your liking. Simply passes the array given. It is the caller's
responsibility to know the array to pass and that the pixel size has been accounted
for already.
If the passed array is smaller than the simulator's image, it will be centered
on the simulator image and added. If the passed array is bigger than the simulator's image, it will
be centered, cropped to fit the simulator's array size, and then added.
Parameters
----------
image : np.ndarray
The array to add to the simulator image.
"""
def __init__(self, image: np.ndarray):
self.array = image
[docs]
def apply(
self, image: np.ndarray, pixel_size: float, mag_factor: float
) -> np.ndarray:
return add_centered_array(base_array=image, other_array=self.array)
def rotate_point(x: float, y: float, angle: float) -> (float, float):
"""
Rotate a point (px, py) about the origin by a given angle in degrees.
Parameters
----------
x : float
The x-coordinate of the point to rotate.
y : float
The y-coordinate of the point to rotate.
angle : float
The angle of rotation in degrees.
Returns
-------
px_rotated, py_rotated : tuple
The rotated point.
Notes
-----
ChatGPT-generated 🥰
"""
# Convert angle from degrees to radians
theta = np.radians(angle)
# Apply rotation
px_rotated = x * np.cos(theta) - y * np.sin(theta)
py_rotated = x * np.sin(theta) + y * np.cos(theta)
return px_rotated, py_rotated
def draw_rotated_rectangle(
shape: tuple[int, int],
center: list[float, float],
extent: list[int, int],
angle: float,
):
"""Generates the coordinate points for a rectangle rotated about the image center.
Parameters
----------
shape : tuple
The shape of the image.
center : list
The center of the rectangle.
extent : list
The height and width of the rectangle.
angle : float
The angle of rotation in degrees.
Returns
-------
rr, cc : tuple
The row and column coordinates of the rectangle.
Notes
-----
ChatGPT-generated 🥰
"""
# Calculate rectangle coordinates before rotation
# follows row,col convention of numpy; y = row, x = col
x0 = center[1] - extent[1] / 2
x1 = center[1] + extent[1] / 2
y0 = center[0] - extent[0] / 2
y1 = center[0] + extent[0] / 2
rect_coords = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]])
# Calculate rotation matrix
theta = np.radians(angle)
c, s = np.cos(theta), np.sin(theta)
rotation = np.array([[c, -s], [s, c]])
# Rotate rectangle coordinates
rotated_coords = np.dot(rect_coords - np.array(center), rotation) + np.array(center)
# Draw rotated rectangle
rr, cc = polygon(rotated_coords[:, 1], rotated_coords[:, 0], shape)
return rr, cc
def add_centered_array(base_array: np.ndarray, other_array: np.ndarray) -> np.ndarray:
"""
Adds a 2D array to another 2D array, centering it over the base array.
If the other array is larger, it will be cropped to fit within the base array.
.. warning:: Will cast the other array to the dtype of the base array.
Parameters
----------
base_array : np.ndarray
The array that the other array will be added to.
other_array : np.ndarray
The array that will be added to the base array.
Returns
-------
np.ndarray
The resulting array after adding the other array centered over the base array.
"""
# Calculate the shapes and start indices
base_shape = np.array(base_array.shape)
other_shape = np.array(other_array.shape).astype(base_array.dtype)
start_indices_base = np.maximum((base_shape - other_shape) // 2, 0)
start_indices_other = np.maximum((other_shape - base_shape) // 2, 0)
# Calculate the end indices
end_indices_base = np.minimum(start_indices_base + other_shape, base_shape)
end_indices_other = np.minimum(start_indices_other + base_shape, other_shape)
# Calculate the slice ranges
base_slice = (
slice(start_indices_base[0], end_indices_base[0]),
slice(start_indices_base[1], end_indices_base[1]),
)
other_slice = (
slice(start_indices_other[0], end_indices_other[0]),
slice(start_indices_other[1], end_indices_other[1]),
)
# Add the sliced arrays
base_array[base_slice] += other_array.astype(base_array.dtype)[other_slice]
return base_array