Registration Module

PyTorch functions for 2D/3D affine image registration, volume warping, and point cloud alignment.

This module provides differentiable implementations for aligning 2D images and 3D volumes using geometric transformations. It includes tools for both intensity-based image registration and geometry-based point cloud registration (ICP).

Main features: - Differentiable 2D and 3D registration (register_affine_2d and register_affine_3d) supporting Rotation, Translation, Scale, and Shear. - Batched volume warping (warp_volume_xy_batched) to apply 2D transforms to 3D stacks efficiently. - Point cloud alignment using a differentiable Iterative Closest Point (icp_torch) implementation. - Utilities for handling pixel-space vs. normalized-space affine matrices.

Author: Antony Vamvakeros

class nDTomo.pytorch.reg_torch.NCC(eps=1e-06)[source]

Bases: object

Zero-Normalized Cross Correlation (robust to brightness changes).

forward(x, y)[source]
nDTomo.pytorch.reg_torch.build_affine_matrix(theta, tx, ty, sx, sy, shx, shy, device, dtype)[source]

Constructs the 3x3 Inverse Affine Matrix (normalized coordinates) combining rotation, translation, scale, and shear.

nDTomo.pytorch.reg_torch.build_affine_matrix_3d(rx, ry, rz, tx, ty, tz, sx, sy, sz, device, dtype)[source]

Constructs the 3x4 Inverse Affine Matrix for 3D volumetric transformations.

Combines Translation, Rotation (Euler angles XYZ), and Scaling. Shear is omitted for simplicity but can be added if needed.

Parameters:

  • rx (torch.Tensor) – Rotation angles around X, Y, and Z axes (in radians).

  • ry (torch.Tensor) – Rotation angles around X, Y, and Z axes (in radians).

  • rz (torch.Tensor) – Rotation angles around X, Y, and Z axes (in radians).

  • tx (torch.Tensor) – Translation in x, y, z (normalized coordinates [-1, 1]).

  • ty (torch.Tensor) – Translation in x, y, z (normalized coordinates [-1, 1]).

  • tz (torch.Tensor) – Translation in x, y, z (normalized coordinates [-1, 1]).

  • sx (torch.Tensor) – Scale factors for x, y, z.

  • sy (torch.Tensor) – Scale factors for x, y, z.

  • sz (torch.Tensor) – Scale factors for x, y, z.

Returns:

(1, 3, 4) tensor representing the top three rows of the 4x4 inverse affine matrix, suitable for 3D F.affine_grid.

Return type:

torch.Tensor

nDTomo.pytorch.reg_torch.icp_torch(A, B, max_iterations=2000, lr=0.01, tolerance=1e-06, verbose=False)[source]

Perform rigid Iterative Closest Point (ICP) registration from point cloud B to A.

Optimizes rotation (R) and translation (t) such that B_aligned = B @ R.T + t fits A.

Parameters:

  • A (torch.Tensor or numpy.ndarray) – Point clouds of shape (N, 3).

  • B (torch.Tensor or numpy.ndarray) – Point clouds of shape (N, 3).

  • max_iterations (int) – Maximum optimization steps.

  • lr (float) – Learning rate.

  • tolerance (float) – Convergence threshold.

Returns:

  • B_aligned (numpy.ndarray) – Aligned point cloud.

  • R (numpy.ndarray) – 3x3 Rotation matrix.

  • t (numpy.ndarray) – Translation vector.

nDTomo.pytorch.reg_torch.mae_loss(x, y)[source]

Mean Absolute Error

nDTomo.pytorch.reg_torch.mse_loss(x, y)[source]

Mean Squared Error

nDTomo.pytorch.reg_torch.normalize_affine_matrix(matrix_pixel, height, width)[source]

Converts a 3x3 affine matrix defined in pixel coordinates (e.g., from skimage or pystackreg) into the normalized coordinate system [-1, 1] required by PyTorch.

Parameters:

  • matrix_pixel (numpy.ndarray) – 3x3 affine matrix in pixel units.

  • height (int) – Dimensions of the image.

  • width (int) – Dimensions of the image.

Returns:

(1, 2, 3) tensor ready for F.affine_grid.

Return type:

torch.Tensor

nDTomo.pytorch.reg_torch.register_affine_2d(ref, mov, order=['rot', 'trans', 'scale', 'shear'], loss_type='ncc', iters=200, lr=0.01, verbose=False, device=None)[source]

Register two 2D images using optimization of affine parameters.

Parameters:

  • ref ((H,W) numpy array or torch tensor)

  • mov ((H,W) numpy array or torch tensor)

  • order (list of str) – Parameters to optimize: ‘rot’, ‘trans’, ‘scale’, ‘shear’.

  • loss_type (str) – Objective function: ‘ncc’, ‘mae’, or ‘mse’.

  • iters (int) – Number of iterations.

  • lr (float) – Learning rate.

Returns:

  • warped_image ((H,W) numpy array)

  • params (dict) – Optimized parameters (theta, tx, ty, sx, sy, etc.)

nDTomo.pytorch.reg_torch.register_affine_3d(ref, mov, order=['rot', 'trans', 'scale'], loss_type='ncc', iters=200, lr=0.01, verbose=False, device=None)[source]

Register two 3D volumes using optimization of affine parameters.

Parameters:

  • ref ((D,H,W) numpy array or torch tensor) – Input volumes.

  • mov ((D,H,W) numpy array or torch tensor) – Input volumes.

  • order (list of str) – Parameters to optimize: ‘rot’, ‘trans’, ‘scale’.

  • loss_type (str) – Objective function: ‘ncc’, ‘mae’, or ‘mse’.

  • iters (int) – Number of iterations.

  • lr (float) – Learning rate.

Returns:

  • warped_vol ((D,H,W) numpy array) – The registered moving volume.

  • params (dict) – Optimized parameters (rx, ry, rz, tx, ty, tz, sx, sy, sz).

nDTomo.pytorch.reg_torch.to_tensor2d(x, device=None, dtype='float32')[source]

Converts a 2D NumPy array or Tensor into a 4D Torch Tensor (1, 1, H, W).

nDTomo.pytorch.reg_torch.to_tensor3d(x, device=None, dtype='float32')[source]

Converts a 3D NumPy array or Tensor into a 5D Torch Tensor (1, 1, D, H, W).

nDTomo.pytorch.reg_torch.warp_volume_xy_batched(volume_np, affine_matrix, is_pixel_space=False, batch_size=32, device=None)[source]

Apply a single 2D affine transform to every Z-slice of a 3D volume using batches.

Parameters:

  • volume_np (numpy.ndarray) – Input 3D volume (Z, H, W).

  • affine_matrix (numpy.ndarray) – 3x3 Affine matrix.

  • is_pixel_space (bool, optional) – If True, assumes affine_matrix is in pixel coordinates (e.g. from Skimage) and normalizes it. If False, assumes it is already PyTorch-compatible [-1, 1].

  • batch_size (int) – Number of slices to process simultaneously.

  • device (str) – ‘cuda’ or ‘cpu’.

Returns:

Transformed 3D volume (Z, H, W).

Return type:

numpy.ndarray