import warnings
import numpy as np
import jax.numpy as jnp
import jax
import mbirjax.bn256 as bn
import mbirjax.preprocess as mjp
def get_2d_ror_mask(recon_shape, *, use_ror_mask=True, crop_radius_pixels=0, crop_radius_fraction=0.0):
"""
Get a binary mask for the region of reconstruction. By default, the mask is an ellipse inscribed in
the edges of the 2D recon_shape[0:2]. The size of this ellipse can be reduced by setting
crop_radius_pixels or crop_radius_fraction, either of which is subtracted from the ellipse axes. Only
one of these can be nonzero. Negative values are clipped to 0.
Args:
recon_shape (tuple): Shape of recon in (rows, columns, slices), or just (rows, columns).
use_ror_mask (default is True):
False:
No mask.
True:
The mask is an ellipse inscribed in the reconstruction volume.
2D array:
Use a custom binary mask. Must have shape recon_shape[:2].
crop_radius_pixels (int): Number of column-pixel-equivalent pixels to subtract from radius.
crop_radius_fraction (float): Fraction to subtract from each axis radius.
Returns:
np.ndarray: Boolean 2D binary mask.
"""
if use_ror_mask is False:
if crop_radius_pixels != 0 and crop_radius_fraction != 0.0:
raise ValueError('crop_radius_pixels and crop_radius_fraction must be zero if use_ror_mask is set to False.')
return np.ones(recon_shape[:2], dtype=bool)
elif use_ror_mask is True:
# Set up a mask to zero out points outside the ROR
if crop_radius_pixels != 0 and crop_radius_fraction != 0.0:
raise ValueError('Only one of crop_radius_pixels and crop_radius_fraction can be nonzero.')
num_recon_rows, num_recon_cols = recon_shape[:2]
row_center = (num_recon_rows - 1) / 2
col_center = (num_recon_cols - 1) / 2
row_radius = row_center - max(int(row_center * crop_radius_fraction), crop_radius_pixels, 0)
col_radius = col_center - max(int(col_center * crop_radius_fraction), crop_radius_pixels, 0)
col_coords = np.arange(num_recon_cols) - col_center
row_coords = np.arange(num_recon_rows) - row_center
coords = np.meshgrid(col_coords, row_coords)
mask = (coords[0] / col_radius) ** 2 + (coords[1] / row_radius) ** 2 <= 1.0
mask = mask[:, :]
return mask
else: # user-provided mask
if crop_radius_pixels != 0 and crop_radius_fraction != 0.0:
raise ValueError('crop_radius_pixels and crop_radius_fraction must be zero if use_ror_mask is a custom array.')
mask = np.asarray(use_ror_mask)
if mask.shape != tuple(recon_shape[:2]):
raise ValueError(
"Custom use_ror_mask must have shape recon_shape[:2]. "
f"Got mask.shape={mask.shape}, expected {tuple(recon_shape[:2])}."
)
if not np.all((mask == 0) | (mask == 1)):
raise ValueError("Custom use_ror_mask must contain only 0s and 1s.")
return mask
def gen_set_of_pixel_partitions(recon_shape, granularity, output_device=None, use_ror_mask=True):
"""
Generates a collection of voxel partitions for an array of specified partition sizes.
This function creates a tuple of randomly generated 2D voxel partitions.
Args:
recon_shape (tuple): Shape of recon in (rows, columns, slices)
granularity (list or tuple): List of num_subsets to use for each partition
output_device (jax device): Device on which to place the output of the partition
use_ror_mask (default is True):
False:
No mask.
True:
The mask is an ellipse inscribed in the reconstruction volume.
2D array:
Use a custom binary mask. Must have shape recon_shape[:2].
Returns:
tuple: A tuple of 2D arrays each representing a partition of voxels into the specified number of subsets.
"""
partitions = []
for num_subsets in granularity:
partition = gen_pixel_partition(recon_shape, num_subsets, use_ror_mask=use_ror_mask)
partitions += [jax.device_put(partition, output_device),]
return partitions
def gen_pixel_partition_grid(recon_shape, num_subsets, use_ror_mask=True):
small_tile_side = np.ceil(np.sqrt(num_subsets)).astype(int)
num_subsets = small_tile_side ** 2
num_small_tiles = [np.ceil(recon_shape[k] / small_tile_side).astype(int) for k in [0, 1]]
single_subset_inds = np.random.permutation(num_subsets).reshape((small_tile_side, small_tile_side))
subset_inds = np.tile(single_subset_inds, num_small_tiles)
subset_inds = subset_inds[:recon_shape[0], :recon_shape[1]]
if use_ror_mask is not False:
ror_mask = get_2d_ror_mask(recon_shape, use_ror_mask=use_ror_mask)
else:
ror_mask = 1
subset_inds = (subset_inds + 1) * ror_mask - 1 # Get a -1 at each location outside the mask, subset_ind at other points
subset_inds = subset_inds.flatten()
num_inds = len(np.where(subset_inds > -1)[0])
if num_subsets > num_inds:
# num_subsets = len(indices)
warning = '\nThe number of partition subsets is greater than the number of pixels in the region of '
warning += 'reconstruction. \nReducing the number of subsets to equal the number of indices.'
warnings.warn(warning)
subset_inds = subset_inds[subset_inds > -1]
return jnp.array(subset_inds).reshape((-1, 1))
flat_inds = []
max_points = 0
min_points = subset_inds.size
nonempty_subsets = np.unique(subset_inds[subset_inds>=0])
for k in nonempty_subsets:
cur_inds = np.where(subset_inds == k)[0]
flat_inds.append(cur_inds) # Get all the indices for each subset
max_points = max(max_points, cur_inds.size)
min_points = min(min_points, cur_inds.size)
extra_point_inds = np.random.randint(min_points, size=(max_points - min_points + 1,))
for k in range(len(nonempty_subsets)):
cur_inds = flat_inds[k]
num_extra_points = max_points - cur_inds.size
if num_extra_points > 0:
extra_subset_inds = (k + 1 + np.arange(num_extra_points, dtype=int)) % len(nonempty_subsets)
new_point_inds = [flat_inds[extra_subset_inds[j]][extra_point_inds[j]] for j in range(num_extra_points)]
flat_inds[k] = np.concatenate((cur_inds, new_point_inds))
flat_inds = np.array(flat_inds)
# Reorganize into subsets, then sort each subset
indices = jnp.array(flat_inds)
return jnp.array(indices)
def gen_pixel_partition(recon_shape, num_subsets, use_ror_mask=True):
"""
Generates a partition of pixel indices into specified number of subsets for use in tomographic reconstruction algorithms.
The function ensures that each subset contains an equal number of pixels, suitable for VCD reconstruction.
Args:
recon_shape (tuple): Shape of recon in (rows, columns, slices)
num_subsets (int): The number of subsets to divide the pixel indices into.
use_ror_mask (default is True):
False:
No mask.
True:
The mask is an ellipse inscribed in the reconstruction volume.
2D array:
Use a custom binary mask. Must have shape recon_shape[:2].
Raises:
ValueError: If the number of subsets specified is greater than the total number of pixels in the grid.
Returns:
jnp.array: A JAX array where each row corresponds to a subset of pixel indices, sorted within each subset.
"""
# Determine the 2D indices within the RoR
num_recon_rows, num_recon_cols = recon_shape[:2]
max_index_val = num_recon_rows * num_recon_cols
indices = np.arange(max_index_val, dtype=np.int32)
# Mask off indices that are outside the region of reconstruction
if use_ror_mask is not False:
mask = get_2d_ror_mask(recon_shape, use_ror_mask=use_ror_mask)
mask = mask.flatten()
indices = indices[mask == 1]
if num_subsets > len(indices):
num_subsets = len(indices)
warning = '\nThe number of partition subsets is greater than the number of pixels in the region of '
warning += 'reconstruction. \nReducing the number of subsets to equal the number of indices.'
warnings.warn(warning)
# A single subset needs no permutation: the subsets are SORTED below, so a shuffle would be
# exactly undone -- but it would consume global np.random state. Skipping it keeps
# full-index "partitions" (gen_full_indices: the Hessian diagonal, the direct-recon init)
# from advancing the RNG, so a restarted recon (init_recon + first_iteration, with
# caller-held partitions) reproduces a continuous run's per-iteration subset
# permutations -- and hence its trajectory -- exactly. The returned partition is
# bit-identical to the permute-then-sort path.
if num_subsets == 1:
return jnp.array(np.sort(indices).reshape(1, -1))
# Determine the number of indices to repeat to make the total number divisible by num_subsets
num_indices_per_subset = int(np.ceil((len(indices) / num_subsets)))
array_size = num_subsets * num_indices_per_subset
num_extra_indices = array_size - len(indices) # Note that this is >=0 since otherwise the ValueError is raised
indices = np.random.permutation(indices)
# Enlarge the array to the desired length by adding random indices that are not in the final subset
num_non_final_indices = (num_subsets - 1) * num_indices_per_subset
extra_indices = np.random.choice(indices[:num_non_final_indices], size=num_extra_indices, replace=False)
indices = np.concatenate((indices, extra_indices))
# Reorganize into subsets, then sort each subset
indices = indices.reshape(num_subsets, indices.size // num_subsets)
indices = jnp.sort(indices, axis=1)
return jnp.array(indices)
def gen_pixel_partition_blue_noise(recon_shape, num_subsets, use_ror_mask=True):
"""
Generates a partition of pixel indices into specified number of subsets for use in tomographic reconstruction algorithms.
The function ensures that each subset contains an equal number of pixels, suitable for VCD reconstruction.
Args:
recon_shape (tuple): Shape of recon in (rows, columns, slices)
num_subsets (int): The number of subsets to divide the pixel indices into.
use_ror_mask (default is True):
False:
No mask.
True:
The mask is an ellipse inscribed in the reconstruction volume.
2D array:
Use a custom binary mask. Must have shape recon_shape[:2].
Raises:
ValueError: If the number of subsets specified is greater than the total number of pixels in the grid.
Returns:
jnp.array: A JAX array where each row corresponds to a subset of pixel indices, sorted within each subset.
"""
pattern = bn.bn256
num_tiles = [np.ceil(recon_shape[k] / pattern.shape[k]).astype(int) for k in [0, 1]]
if use_ror_mask is not False:
ror_mask = get_2d_ror_mask(recon_shape, use_ror_mask=use_ror_mask)
else:
ror_mask = 1
single_subset_inds = np.floor(pattern / (2**16 / num_subsets)).astype(int)
# Repeat each bn subset to do the tiling
subset_inds = np.tile(single_subset_inds, num_tiles)
subset_inds = subset_inds[:recon_shape[0], :recon_shape[1]]
subset_inds = (subset_inds + 1) * ror_mask - 1 # Get a -1 at each location outside the mask, subset_ind at other points
subset_inds = subset_inds.flatten()
num_valid_inds = np.sum(subset_inds >= 0)
if num_subsets > num_valid_inds:
return gen_pixel_partition(recon_shape, num_subsets, use_ror_mask=use_ror_mask)
flat_inds = []
max_points = 0
min_points = subset_inds.size
for k in range(num_subsets):
cur_inds = np.where(subset_inds == k)[0]
flat_inds.append(cur_inds) # Get all the indices for each subset
max_points = max(max_points, cur_inds.size)
min_points = min(min_points, cur_inds.size)
if min_points == 0:
return gen_pixel_partition(recon_shape, num_subsets, use_ror_mask=use_ror_mask)
extra_point_inds = np.random.randint(low=0, high=min_points, size=(max_points - min_points + 1,))
for k in range(num_subsets):
cur_inds = flat_inds[k]
num_extra_points = max_points - cur_inds.size
if num_extra_points > 0:
extra_subset_inds = (k + 1 + np.arange(num_extra_points, dtype=int)) % num_subsets
new_point_inds = [flat_inds[extra_subset_inds[j]][extra_point_inds[j]] for j in range(num_extra_points)]
flat_inds[k] = np.concatenate((cur_inds, new_point_inds))
flat_inds = jnp.array(flat_inds)
flat_inds = jnp.sort(flat_inds, axis=1)
return flat_inds
def gen_partition_sequence(partition_sequence, max_iterations):
"""
Generates a sequence of voxel partitions of the specified length by extending the sequence
with the last element if necessary.
"""
# Get sequence from params and convert it to a np array
partition_sequence = jnp.array(partition_sequence)
# Check if the sequence needs to be extended
current_length = partition_sequence.size
if max_iterations > current_length:
# Calculate the number of additional elements needed
extra_elements_needed = max_iterations - current_length
# Get the last element of the array
last_element = partition_sequence[-1]
# Create an array of the last element repeated the necessary number of times
extension_array = np.full(extra_elements_needed, last_element)
# Concatenate the original array with the extension array
extended_partition_sequence = np.concatenate((partition_sequence, extension_array))
else:
# If no extension is needed, slice the original array to the desired length
extended_partition_sequence = partition_sequence[:max_iterations]
return extended_partition_sequence
def gen_full_indices(recon_shape, use_ror_mask=True):
"""
Generates a full array of voxels in the region of reconstruction.
This is useful for computing forward projections.
"""
partition = gen_pixel_partition(recon_shape, num_subsets=1, use_ror_mask=use_ror_mask)
full_indices = partition[0]
return full_indices
[docs]
def gen_weights_mar(ct_model, sinogram, init_recon=None, metal_threshold=None, beta=1.0, gamma=3.0):
"""
Generates the weights used for reducing metal artifacts in MBIR reconstruction.
This function computes sinogram weights that help to reduce metal artifacts.
More specifically, it computes weights with the form:
weights = exp( -(sinogram/beta) * ( 1 + gamma * delta(metal) ) )
delta(metal) denotes a binary mask indicating the sino entries that contain projections of metal.
Providing ``init_recon`` yields better metal artifact reduction.
If not provided, the metal segmentation is generated directly from the sinogram.
Args:
sinogram (jax array): 3D jax array containing sinogram with shape (num_views, num_det_rows, num_det_channels).
init_recon (jax array, optional): An initial reconstruction used to identify metal voxels. If not provided, Otsu's method is used to directly segment sinogram into metal regions.
metal_threshold (float, optional): Values in ``init_recon`` above ``metal_threshold`` are classified as metal. If not provided, Otsu's method is used to segment ``init_recon``.
beta (float, optional): Scalar value in range :math:`>0`.
A larger ``beta`` improves the noise uniformity, but too large a value may increase the overall noise level.
gamma (float, optional): Scalar value in range :math:`>=0`.
A larger ``gamma`` reduces the weight of sinogram entries with metal, but too large a value may reduce image quality inside the metal regions.
Returns:
(jax array): Weights used in mbircone reconstruction, with the same array shape as ``sinogram``
"""
# The masks below are element-wise comparisons, so they inherit the sharding of their source
# (the sinogram or init_recon) rather than being pinned to a single device -- that keeps the
# init_recon=None path sharding-transparent (delta_metal follows the sinogram, so the final
# element-wise weights are sharded with no gather). NOTE: the init_recon path still materializes
# single-device intermediates -- forward_project here uses the default gathered output, and
# multi_threshold_otsu may gather -- so it is not yet memory-bounded for very large sharded
# problems; sharding that path is the mar/preprocessing follow-on.
# If init_recon is not provided, then identify the distorted sino entries with Otsu's thresholding method.
if init_recon is None:
print("init_recon is not provided. Automatically determine distorted sinogram entries with Otsu's method.")
# assuming three categories: metal, non_metal, and background.
[bk_thresh_sino, metal_thresh_sino] = mjp.multi_threshold_otsu(sinogram, classes=3)
print("Distorted sinogram threshold = ", metal_thresh_sino)
delta_metal = (sinogram > metal_thresh_sino).astype(jnp.float32)
# If init_recon is provided, identify the distorted sino entries by forward projecting init_recon.
else:
if metal_threshold is None:
print("Metal_threshold calculated with Otsu's method.")
# assuming three categories: metal, non_metal, and background.
[bk_threshold, metal_threshold] = mjp.multi_threshold_otsu(init_recon, classes=3)
print("metal_threshold = ", metal_threshold)
# Identify metal voxels
metal_mask = (init_recon > metal_threshold).astype(jnp.float32)
# Forward project metal mask to generate a sinogram mask
metal_mask_projected = ct_model.forward_project(metal_mask)
# metal mask in the sinogram domain, where 1 means a distorted sino entry, and 0 else.
delta_metal = (metal_mask_projected > 0.0).astype(jnp.float32)
# weights for undistorted sino entries
weights = jnp.exp(-sinogram*(1+gamma*delta_metal)/beta)
return weights
[docs]
def gen_weights(sinogram, weight_type, ct_model=None):
"""
Compute optional weights used in MBIR reconstruction based on the noise model.
The weights should be proportional to the inverse variance of the noise for each sinogram entry.
They can be used to improve reconstruction quality.
The result is returned where the input lives, so it feeds reconstruction without landing a large
array on a single device:
- a host (NumPy) sinogram yields host weights (reconstruction streams them to the devices
shard-by-shard);
- a JAX array yields JAX weights inheriting its sharding (a view-sharded sinogram gives
view-sharded weights, with no cross-device communication or gather).
To build the weights already distributed across a multi-device reconstruction -- avoiding any
single-device copy of a large host sinogram -- pass the model as ``ct_model``: the sinogram is
placed in the model's view-sharded device form first, then weighted per shard.
Args:
sinogram (ndarray or jax.Array): A 3D array of shape (num_views, num_det_rows, num_det_channels)
representing the sinogram.
weight_type (str): The type of noise model to use for weighting. Must be one of:
- 'unweighted': Use uniform weights (all ones).
- 'transmission': Use exponential decay, `exp(-sinogram)`.
- 'transmission_root': Use square-root decay, `exp(-sinogram / 2)`.
- 'emission': Use reciprocal decay, `1 / (abs(sinogram) + 0.1)`.
ct_model (TomographyModel, optional): If given, distribute the sinogram in this model's
sinogram device form (view-sharded across its devices) before weighting, so the weights
are returned sharded and ready for reconstruction with no single-device copy. Defaults to
None (the result preserves the input's residence, as described above).
Returns:
ndarray or jax.Array: A 3D array of weights with the same shape as the input sinogram, on the
host or sharded across devices to match the input (or ``ct_model``).
Raises:
Exception: If `weight_type` is not one of the supported options.
Note:
For transmission noise models, sinogram values should not be excessively large (e.g., > 5),
as this corresponds to near-zero transmission, which is not physically meaningful in typical X-ray imaging.
Example:
>>> sinogram = jnp.ones((180, 64, 128))
>>> weights = gen_weights(sinogram, weight_type='transmission_root')
>>> weights.shape
(180, 64, 128)
"""
# Optionally place the sinogram in the model's view-sharded device form first, so the weights are
# built per shard (no single-device copy of a large host sinogram).
if ct_model is not None:
sinogram = ct_model._shard_sinogram(sinogram)
# The weights are an element-wise function of the sinogram, computed in one pass with the input's
# OWN array module so the result stays where the input is: a host (NumPy) sinogram yields host
# weights, a JAX array yields JAX weights inheriting its sharding (each shard processed in place --
# element-wise ops need no cross-device communication). This keeps a large host sinogram OFF the
# devices until reconstruction streams it shard-by-shard, and never lands the whole array on one
# device. On a single device the peak is ~2x the sinogram (input + output); shard the sinogram
# (or pass ct_model) when it is too large for one device.
xp = jnp if isinstance(sinogram, jax.Array) else np
if weight_type == 'unweighted':
weights = xp.ones_like(sinogram) # ones_like (not ones) preserves the input's sharding
elif weight_type == 'transmission':
weights = xp.exp(-sinogram)
elif weight_type == 'transmission_root':
weights = xp.exp(-sinogram / 2)
elif weight_type == 'emission':
weights = 1.0 / (xp.absolute(sinogram) + 0.1)
else:
raise Exception("gen_weights: undefined weight_type {}".format(weight_type))
return weights
