Source code for sphero_vem.registration.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.
"""

import numpy as np
from skimage import measure


[docs] def border_mask(image: np.ndarray) -> np.ndarray: """ 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 ------- np.ndarray Boolean mask where True is border and False is image content. """ masked = (image == 0) | (image == 255) labeled = measure.label(masked, connectivity=1) # Collect component ids that touch any border touching = set() touching.update(np.unique(labeled[0, :])) touching.update(np.unique(labeled[image.shape[-2] - 1, :])) touching.update(np.unique(labeled[:, 0])) touching.update(np.unique(labeled[:, image.shape[-1] - 1])) # Remove background label touching.discard(0) if not touching: return np.zeros_like(image, dtype=bool) border_mask_result = np.isin(labeled, list(touching)) return border_mask_result
[docs] def integral_image(mask: np.ndarray) -> np.ndarray: """ 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 ------- np.ndarray Padded integral image of mask, with size (H + 1, W + 1) """ H, W = mask.shape ii_padded = np.zeros((H + 1, W + 1), dtype=mask.dtype) ii_padded[1:, 1:] = mask return ii_padded.cumsum(0).cumsum(1)
[docs] def rect_sum(ii_padded: np.ndarray, up: int, down: int, left: int, right: int) -> int: """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 ------- int Sum of values in the specified rectangle. """ return ( ii_padded[down, right] - ii_padded[up, right] - ii_padded[down, left] + ii_padded[up, left] )
def _area(box: tuple[int, int, int, int]) -> int: """Compute the pixel area of a bounding box given as (up, down, left, right). Parameters ---------- box : tuple[int, int, int, int] Bounding box as (up, down, left, right) indices. Returns ------- int Area in pixels: ``(down - up) * (right - left)``. """ return (box[1] - box[0]) * (box[3] - box[2])
[docs] def refine_crop( ii: np.ndarray, up: int, down: int, left: int, right: int, rng: int | None = None ) -> tuple[int, int, int, int]: """ 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, down, left, 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 : int The updated indices. """ rng = np.random.default_rng(rng) improved = True while improved: improved = False def try_expand_up(): nonlocal up, improved if up > 0 and rect_sum(ii, up - 1, down, left, right) == 0: up -= 1 improved = True def try_expand_down(): nonlocal down, improved if down < ii.shape[0] - 1 and rect_sum(ii, up, down + 1, left, right) == 0: down += 1 improved = True def try_expand_left(): nonlocal left, improved if left > 0 and rect_sum(ii, up, down, left - 1, right) == 0: left -= 1 improved = True def try_expand_right(): nonlocal right, improved if right < ii.shape[1] - 1 and rect_sum(ii, up, down, left, right + 1) == 0: right += 1 improved = True moves = [try_expand_up, try_expand_down, try_expand_left, try_expand_right] rng.shuffle(moves) for move in moves: move() return up, down, left, right
[docs] def refine_crop_multistart( ii: np.ndarray, seed_box: tuple[int, int, int, int], n_restarts: int = 10, jitter: int = 10, rng: int | None = None, ) -> tuple[int, int, int, int]: """ 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) : tuple[int] The best indices found by the algorithm. """ rng = np.random.default_rng(rng) def jitter_h(idx: int) -> int: """Add jitter to an integer index in height (up, down)""" return min(ii.shape[0] - 1, idx + rng.integers(-jitter, jitter + 1)) def jitter_w(idx: int) -> int: """Add jitter to an integer index in width (left, right)""" return min(ii.shape[1] - 1, idx + rng.integers(-jitter, jitter + 1)) best_box = refine_crop(ii, *seed_box, rng=rng) best_area = _area(best_box) up_0, down_0, left_0, right_0 = seed_box for _ in range(n_restarts): up = max(0, jitter_h(up_0)) down = max(up, jitter_h(down_0)) left = max(0, jitter_w(left_0)) right = max(left, jitter_w(right_0)) # Check if box is valid if rect_sum(ii, up, down, left, right) != 0: continue cand_crop = refine_crop(ii, up, down, left, right, rng=rng) cand_area = _area(cand_crop) if cand_area > best_area: best_box, best_area = cand_crop, cand_area return best_box
[docs] def find_border_crop( image: np.ndarray, pix_stride: int = 20, n_restarts: int = 10, jitter: int = 10, rng: int | None = None, ) -> tuple[int, int, int, int]: """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 ------- tuple[int, int, int, int] Crop indices as (up, down, left, right). """ masked = border_mask(image) ii = integral_image(masked) seed_box = rough_crop_search(ii, pix_stride) crop_idx = refine_crop_multistart( ii, seed_box, n_restarts=n_restarts, jitter=jitter, rng=rng ) return crop_idx