Registration (sphero_vem.registration)#

Intensity-based pairwise affine registration of consecutive slices, multi-resolution PyTorch optimization, and border crop detection.

Core#

Core registration pipeline for volume stacks.

This module contains the configuration, pairwise registration algorithm, and the main stack registration pipeline with optional border cropping.

class sphero_vem.registration.core.RegistrationConfig(root_path, src_path, dst_group='images/registered', pyramid_levels=4, pyramid_factors=2, pyramid_epochs=300, learning_rate=<factory>, loss_type='ncc', loss_kwargs=<factory>, optimizer='Adam', max_pairs=None, progress_steps=False, verbose=True, early_stopping=True, stop_window=15, stop_tol=1e-05, transformation='similarity', init_std=0.001, scaling=True, shear=True, regularization_param=0.5, crop_borders=True, crop_safety_margin=5, crop_stride=20, crop_restarts=10, n_workers=4)[source]#

Bases: BaseConfig

Configuration for pairwise affine registration of a Zarr volume stack.

Parameters:

  • root_path (Path) – Root directory of the Zarr archive.

  • src_path (str) – Path within the archive to the source image array (e.g. "images/50-10-10").

  • dst_group (str, optional) – Group path under which registered slices are written. Default is "images/registered".

  • pyramid_levels (int, optional) – Number of resolution levels in the multi-scale pyramid. Default is 4.

  • pyramid_factors (int | list[int], optional) – Downsampling factor(s) between pyramid levels. An integer is expanded to [factor**(levels-1), ..., factor**0]. Default is 2.

  • pyramid_epochs (int | list[int], optional) – Training epochs per pyramid level. A scalar is broadcast to all levels. Default is 300.

  • learning_rate (float | list[float], optional) – Learning rate per pyramid level. Default is [1e-3, 1e-3, 5e-4, 1e-4].

  • loss_type (str, optional) – Name of the image similarity loss (see LossDispatcher). Default is "ncc".

  • loss_kwargs (dict, optional) – Extra keyword arguments forwarded to the loss function. Default is {}.

  • optimizer (str, optional) – Optimizer class name (currently only "Adam" is used). Default is "Adam".

  • max_pairs (int | None, optional) – Maximum number of consecutive slice pairs to register. If None, all pairs in the stack are used. Default is None.

  • progress_steps (bool, optional) – Show per-step tqdm progress bars. Default is False.

  • verbose (bool, optional) – Print high-level progress messages. Default is True.

  • early_stopping (bool, optional) – Enable early stopping when the loss plateaus. Default is True.

  • stop_window (int, optional) – Rolling window size for early-stopping fluctuation check. Default 15.

  • stop_tol (float, optional) – Fluctuation tolerance below which training stops. Default is 1e-5.

  • transformation (str, optional) – Transformation family: "similarity", "rigid", or "affine". Default is "similarity".

  • init_std (float, optional) – Standard deviation for random initialization of transform parameters. Default is 0.001.

  • scaling (bool, optional) – Enable per-axis scaling when transformation="affine". Default True.

  • shear (bool, optional) – Enable shear when transformation="affine". Default is True.

  • regularization_param (float, optional) – Weight for L2 regularization on non-translation parameters. Default is 0.5.

  • crop_borders (bool, optional) – Apply border cropping after registration. Default is True.

  • crop_safety_margin (int, optional) – Extra pixels added to the most restrictive per-slice crop box. Default is 5.

  • crop_stride (int, optional) – Pixel stride for the coarse crop search. Default is 20.

  • crop_restarts (int, optional) – Number of random restarts for the crop refinement. Default is 10.

  • n_workers (int, optional) – Number of Dask threads for writing the output array. Default is 4.

class sphero_vem.registration.core.ImageTransform(config, delta_q_init=None)[source]#

Bases: Module

Learnable 2D affine transform module for image registration.

Parameterizes the transformation as a deviation from the identity so that all learned parameters are zero-centred at initialization.

Parameters:

  • config (RegistrationConfig) – Registration configuration. Controls which transform components (scaling, shear) are active via transformation, scaling, and shear fields, and provides the target device.

  • delta_q_init (torch.Tensor | None, optional) – Initial parameter deviation tensor of shape (1, 6). If None, a zero tensor is used.

forward(img)[source]#

Apply the current learned transform to img.

Parameters:

img (torch.Tensor) – Input image tensor of shape (N, C, H, W).

Returns:

Warped image tensor of the same shape as img.

Return type:

torch.Tensor

sphero_vem.registration.core.register_image_pair(fixed_img, moving_img, config)[source]#

Register a moving image to a fixed image using multi-scale gradient descent.

Optimizes a 2D affine transform at each pyramid level, passing the learned parameters as initialization to the next (finer) level.

Parameters:

  • fixed_img (torch.Tensor) – Fixed reference image of shape (1, 1, H, W).

  • moving_img (torch.Tensor) – Moving image to be aligned, same shape as fixed_img.

  • config (RegistrationConfig) – Registration configuration including pyramid settings, optimizer parameters, loss type, and device.

Return type:

tuple[Tensor, list[float]]

Returns:

  • q_final (torch.Tensor) – Final affine parameter vector of shape (1, 6), as absolute parameters (deviation + identity).

  • full_loss_history (list[list]) – Per-step loss records as [level, loss_value] pairs across all pyramid levels.

sphero_vem.registration.core.register_stack(config)[source]#

Register volume stack from zarr archive.

This function performs registration in two phases:

  1. Sequential affine calculation: Register consecutive image pairs and write to temporary zarr.

  2. Post-processing: Optional border cropping and writing final output.

Parameters:

config (RegistrationConfig) – Configuration object with all registration and processing parameters

Return type:

None

Transforms#

Affine transform math for 2D image registration.

This module contains pure transform operations: building affine matrices from parameters, warping images, and composing sequential transforms. It has no I/O or zarr dependencies — only torch.

class sphero_vem.registration.transforms.TransformType(*values)[source]#

Bases: Enum

Cropping#

Border detection and crop optimization for registered images.

After affine registration, warped images contain black borders (and sometimes white resampling artifacts at value 255). This module provides functions to detect those borders and find the largest border-free rectangular crop across all slices of a volume.

sphero_vem.registration.cropping.border_mask(image)[source]#

Return a boolean mask of elements that belong to components touching the image border.

Parameters:

image (np.ndarray) – The image is assumed to be a np.uint8 grayscale image, and border pixels are those with values 0 or 255 (255 is present due to resampling artifacts). Interior islands with these values are excluded.

Returns:

Boolean mask where True is border and False is image content.

Return type:

np.ndarray

sphero_vem.registration.cropping.integral_image(mask)[source]#

Computes the integral image (summed-area table) of a mask.

The output is padded with a 1-pixel border of zeros on the top and left to allow for easier calculation.

Parameters:

mask (np.ndarray) – Boolean mask where True is border and False is image content, of size (H, W)

Returns:

Padded integral image of mask, with size (H + 1, W + 1)

Return type:

np.ndarray

sphero_vem.registration.cropping.rect_sum(ii_padded, up, down, left, right)[source]#

Compute the sum of values in a rectangular region using an integral image.

The region spans rows [up, down) and columns [left, right).

Parameters:

  • ii_padded (numpy.ndarray) – Padded integral image of shape (H+1, W+1), as produced by integral_image.

  • up (int) – Top row index (inclusive).

  • down (int) – Bottom row index (exclusive).

  • left (int) – Left column index (inclusive).

  • right (int) – Right column index (exclusive).

Returns:

Sum of values in the specified rectangle.

Return type:

int

Return the indices for a clean square crop of the original image.

This function performs a rough search by shrinking a square box and checking if the integral image is 0.

Parameters:

  • ii (np.ndarray) – A padded integral image, such as that generated by integral_image.

  • pix_stride (int) – The stride of the shrinking search. Default is 20.

Returns:

(up, down, left, right) – The indices for the largest clean box found with the given stride.

Return type:

tuple[int]

sphero_vem.registration.cropping.refine_crop(ii, up, down, left, right, rng=None)[source]#

Randomized hill climb for finding the largest border-free box in an image.

Each iteration shuffles the order of side expansions, and expands only if the rectangle stays clean.

Parameters:

  • ii (np.ndarray) – An integral image obtained from a masked array where 0=content and 1=border.

  • up (int) – Initial indices for the search. They must already index a clean box.

  • down (int) – Initial indices for the search. They must already index a clean box.

  • left (int) – Initial indices for the search. They must already index a clean box.

  • right (int) – Initial indices for the search. They must already index a clean box.

  • rng (int | None) – Optional seed for random number generator. Default if None

Returns:

up, down, left, right – The updated indices.

Return type:

int

sphero_vem.registration.cropping.refine_crop_multistart(ii, seed_box, n_restarts=10, jitter=10, rng=None)[source]#

Find the crop in an integral image that is free of border pixels with multiple restarts.

Multiple restarts are used to reduce the chances of being stuck in a local minimum. At each restart the clean starting seed_box is jittered to have a different start. If no rng seed is passed, each iteration will have a different order of exploration.

Parameters:

  • ii (np.ndarray) – An integral image obtained from a masked array where 0=content and 1=border. The integral image should be padded with one row/column of 0 at the top/left, such as those produced by integral_image.

  • seed_box (tuple[int, int, int, int]) – A tuple of indices for a clean crop that will be used to initialize the search. Indices order are (up, down, left, right).

  • n_restarts (int) – The number of restarts for the multistart approach. Default is 10.

  • jitter (int) – Range of the random jitter added to seed_box at each restart. If jittering produces an invalid starting crop (i.e. one that contains border pixels), it will be discarded and a new one will be generated. Default is 10.

  • rng (int | None) – Optional seed for random number generator. This will be used for both the jitter and passed internally to refine_crop. Default if None.

Returns:

(up, down, left, right) – The best indices found by the algorithm.

Return type:

tuple[int]

sphero_vem.registration.cropping.find_border_crop(image, pix_stride=20, n_restarts=10, jitter=10, rng=None)[source]#

Find the largest border-free rectangular crop of a registered image.

Black (0) and full-white (255) pixels arising from affine warping are treated as border artifacts and excluded from the crop region.

Parameters:

  • image (numpy.ndarray) – 2-D uint8 image containing border artifacts.

  • pix_stride (int, optional) – Pixel stride for the initial coarse crop search. Default is 20.

  • n_restarts (int, optional) – Number of random restarts for the hill-climb refinement. Default is 10.

  • jitter (int, optional) – Jitter range (pixels) applied to the seed box at each restart. Default is 10.

  • rng (int | None, optional) – Seed for the random number generator. Default is None.

Returns:

Crop indices as (up, down, left, right).

Return type:

tuple[int, int, int, int]