Utilities Module

Utility functions for patch-based sampling, index generation, total variation (TV) regularization, and structural similarity (SSIM) loss in 2D and 3D tensor volumes. These are commonly used in image reconstruction, inverse problems, and spectral imaging applications.

Author: Antony Vamvakeros

Contents: - Index generation utilities:

  • generate_indices: Draws (row, col) indices from uniform, normal, or Sobol distributions.

  • generate_sobol_indices_batch: Batch generation of Sobol-distributed indices.

  • draw_valid_indices: Wrapper for drawing valid indices using masks and custom sampling.

  • filter_patch_indices: Filters mask into valid patch locations.

  • calc_patches_indices: Extracts tensor patches based on index positions.

  • initialize_counter, update_counter: Track pixel utilization in sampling-based methods.

  • Normalization:

    • denormalize: Converts normalized parameter values [0, 1] back to their physical scale.

  • Total Variation (TV) regularization:

    • tv_spatial: TV along spatial dimensions (y, x) of a 2D image.

    • tv_spectral: TV along spectral (channel) dimension.

    • tv_3d_spectral: TV across both spectral and spatial dimensions (channel, y, x).

    • tv_3d: TV in full 3D volumes (z, y, x), supporting 5D tensors (B, C, D, H, W).

  • Structural Similarity (SSIM) loss functions:

    • SSIM2D: Perceptual similarity loss for 2D images based on luminance, contrast, and structure.

    • SSIM3DLoss: SSIM loss adapted for 3D volumes using a uniform kernel.

Notes: All functions are implemented with PyTorch and designed to run on either CPU or CUDA. The TV and SSIM losses are differentiable and can be integrated directly into model training loops for regularization or perceptual fidelity in reconstruction tasks.

class nDTomo.pytorch.utils_torch.SSIM2D(window_size=11, sigma=1.5, C1=0.0001, C2=0.0009, device='cuda')[source]

Bases: object

A PyTorch module for computing the Structural Similarity Index (SSIM) between two 2D images. SSIM is a perceptual metric that captures similarity in terms of luminance, contrast, and structure.

This implementation supports single-channel 2D images (e.g. grayscale) and uses a 2D Gaussian filter window for local statistics computation.

Parameters:

  • window_size (int, optional) – Size of the Gaussian filter window. Default is 11.

  • sigma (float, optional) – Standard deviation of the Gaussian kernel. Default is 1.5.

  • C1 (float, optional) – Stabilizing constant for luminance term. Default is (0.01)^2.

  • C2 (float, optional) – Stabilizing constant for contrast term. Default is (0.03)^2.

  • device (str, optional) – Device to store the Gaussian window tensor. Default is ‘cuda’.

  • Forward

  • -------

  • img1 (torch.Tensor) – First image tensor of shape (N, 1, H, W).

  • img2 (torch.Tensor) – Second image tensor of shape (N, 1, H, W).

Returns:

loss – A scalar tensor representing 1 - SSIM(img1, img2). Suitable for use as a loss function.

Return type:

torch.Tensor

Notes

  • The output is 1 - SSIM, so this module can be used directly as a loss function in optimization.

  • Assumes input tensors are normalized to [0, 1].

  • Only supports single-channel inputs (channel = 1).

create_window(window_size, sigma)[source]
forward(img1, img2)[source]
static gaussian(window_size, sigma)[source]
class nDTomo.pytorch.utils_torch.SSIM3DLoss(window_size=11)[source]

Bases: object

Computes the Structural Similarity Index (SSIM) loss between two 3D volumes. SSIM is a perceptual similarity metric that evaluates structural fidelity based on local statistics (mean, variance, and covariance).

This implementation uses a uniform 3D window (not Gaussian) to approximate local statistics, and is suitable for comparing volumes in tasks such as CT reconstruction, MRI denoising, or 3D image synthesis.

Parameters:

  • window_size (int, optional) – Size of the cubic window used for computing local means and variances. Default is 11.

  • Forward

  • -------

  • x (torch.Tensor) – First input volume of shape (D, H, W).

  • y (torch.Tensor) – Second input volume of shape (D, H, W).

Returns:

loss – Scalar tensor representing 1 - SSIM(x, y), suitable for optimization as a loss function.

Return type:

torch.Tensor

Notes

  • Input volumes are expected to be normalized to [0, 1].

  • Adds singleton batch and channel dimensions internally.

  • Assumes single-channel input; multi-channel support would require modification.

  • The window is uniform (box filter) rather than Gaussian for simplicity.

create_3D_window(window_size)[source]
forward(x, y)[source]
nDTomo.pytorch.utils_torch.calc_patches_indices(indices, tensor, patch_size, use_middle=False)[source]

Calculate the indices of the patches to be selected based on the provided indices. Users can choose to use the middle of the patch or the top-left as the reference point.

Parameters: indices (list of tuples): A list of tuples containing the starting indices of the patches. tensor (torch.Tensor): The tensor from which the patches are to be selected. patch_size (int): The size of the patches. use_middle (bool): If True, use the middle pixel as the reference. If False, use the top-left.

Returns: torch.Tensor: A tensor containing the selected patches.

nDTomo.pytorch.utils_torch.denormalize(param, param_name, param_min, param_max, peak_number=None)[source]

Denormalize the parameter from [0, 1] to its original range.

Parameters: param (float): The parameter to be denormalized. param_name (str): The name of the parameter. param_min (dict): A dictionary containing the minimum values of the parameters. param_max (dict): A dictionary containing the maximum values of the parameters. peak_number (int, optional): The peak number for the parameter. Default is None.

Returns: float: The denormalized parameter.

nDTomo.pytorch.utils_torch.draw_valid_indices(rows, cols, num_indices, patch_size, distribution_type='normal', batched=False, std_dev=3, mask=None, device='cuda')[source]

Draw valid indices from a specified distribution.

Parameters: rows (int): The total number of rows. cols (int): The total number of columns. num_indices (int): The number of indices to draw. patch_size (int): The size of the square patch. distribution_type (str): The type of distribution to draw from (‘normal’, ‘uniform’ or ‘Sobol’). std_dev (float): The standard deviation for the normal distribution. mask (torch.Tensor): A mask to apply to the indices. device (torch.device): The device on which to generate the indices.

Returns: list: A list of valid indices.

nDTomo.pytorch.utils_torch.filter_patch_indices(mask, patch_size)[source]

Divide a binary mask into non-overlapping patches and return the top-left coordinates of patches that contain any non-zero elements.

Parameters:

  • mask (torch.Tensor) – A 2D binary tensor indicating valid regions (non-zero entries).

  • patch_size (int) – Size of the square patch to extract.

Returns:

patch_indices – A list of (row, col) coordinates indicating the top-left corners of valid patches. The list is randomly shuffled.

Return type:

list of tuple[int, int]

nDTomo.pytorch.utils_torch.generate_indices(rows, cols, num_indices, patch_size, distribution_type='uniform', std=3, mask=None, device='cuda', batched=False)[source]

Generate a specified number of valid (row, col) indices from a 2D grid, optionally constrained by a mask and following a given sampling distribution.

Parameters:

  • rows (int) – Number of rows in the 2D grid.

  • cols (int) – Number of columns in the 2D grid.

  • num_indices (int) – Total number of valid (row, col) pairs to generate.

  • patch_size (int) – Size of the patch, used only for Sobol sampling mode.

  • distribution_type (str, optional) – Sampling distribution to use: ‘uniform’, ‘normal’, or ‘Sobol’. Default is ‘uniform’.

  • std (float, optional) – Standard deviation for ‘normal’ distribution sampling. Ignored for other distributions. Default is 3.

  • mask (torch.Tensor or None, optional) – A binary mask (shape: [rows, cols]) where only non-zero locations are considered valid. Default is None.

  • device (str, optional) – Device to perform computations on (‘cuda’ or ‘cpu’). Default is ‘cuda’.

  • batched (bool, optional) – If True, samples are drawn in batches to speed up sampling. Default is False.

Returns:

indices – A list of (row, col) tuples representing valid sampled locations, possibly mask-constrained.

Return type:

list of tuple[int, int]

nDTomo.pytorch.utils_torch.generate_sobol_indices_batch(rows, cols, batch_size, patch_size, device)[source]

Generate a batch of Sobol indices for a given dimension and batch size, considering the patch size to avoid boundary issues.

Parameters: rows (int): The total number of rows in the dataset. cols (int): The total number of columns in the dataset. batch_size (int): The number of indices to generate. patch_size (int): The size of the patch to consider for boundary adjustment. device (torch.device): The device to use for generating indices.

Returns: tuple: Two tensors containing the row indices and column indices.

nDTomo.pytorch.utils_torch.initialize_counter(rows, cols)[source]

Initialize a zero-filled counter matrix of the same size as the input matrix using PyTorch.

Parameters: - matrix (2D torch tensor): Matrix whose dimensions will be used to create the counter.

Returns: - 2D torch tensor: Initialized counter matrix.

nDTomo.pytorch.utils_torch.tv_3d(x, isotropic=True, epsilon=1e-06)[source]

Compute Total Variation (TV) gradient for 3D volumes along all three axes (x, y, z).

Parameters:

  • x – Input volume of shape (1, nch, D, H, W).

  • isotropic – If True, use isotropic TV. If False, use anisotropic TV.

  • epsilon – Small value to prevent division by zero (for isotropic TV).

Returns:

TV gradient of the same shape as x.

nDTomo.pytorch.utils_torch.tv_3d_spectral(x, isotropic=True, epsilon=1e-06)[source]

Compute Total Variation (TV) gradient along spectral and spatial dimensions (nch, y, x).

Parameters:

  • x – Input volume of shape (1, nch, npix, npix).

  • isotropic – If True, use isotropic TV. If False, use anisotropic TV.

  • epsilon – Small value to prevent division by zero (for isotropic TV).

Returns:

TV gradient of the same shape as x.

nDTomo.pytorch.utils_torch.tv_spatial(x, isotropic=True, epsilon=1e-06)[source]

Compute Total Variation (TV) gradient along spatial dimensions (y, x).

Parameters:

  • x – Input volume of shape (1, nch, npix, npix).

  • isotropic – If True, use isotropic TV. If False, use anisotropic TV.

  • epsilon – Small value to prevent division by zero (for isotropic TV).

Returns:

TV gradient of the same shape as x.

nDTomo.pytorch.utils_torch.tv_spectral(x, isotropic=True, epsilon=1e-06)[source]

Compute Total Variation (TV) gradient along the spectral dimension (nch).

Parameters:

  • x – Input volume of shape (1, nch, npix, npix).

  • isotropic – If True, use isotropic TV. If False, use anisotropic TV.

  • epsilon – Small value to prevent division by zero (for isotropic TV).

Returns:

TV gradient of the same shape as x.

nDTomo.pytorch.utils_torch.update_counter(counter, indices, patch_size=(1, 1))[source]

Update the pixel utilization counter based on the middle indices of extracted patches or pixels using PyTorch.

Parameters: - counter (2D torch tensor): Counter matrix to be updated. - indices (list of tuples): Indices of the top-left corners of the patches or pixels. - patch_size (tuple): Size (height, width) of the patches or single pixel (default is 1x1 for pixels).