from __future__ import annotations
import json
import logging
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
from ewokscore import Task
from ewokscore.model import BaseInputModel, BaseOutputModel
from pydantic import Field
from scipy.signal import correlate, correlation_lags
from tifffile import imwrite
from tomoscan.esrf.volume import MultiTIFFVolume, TIFFVolume
from tomoscan.volumebase import VolumeBase
logger = logging.getLogger(__name__)
Z_AXIS = 0
X_AXIS = 1
Y_AXIS = 2
AXIS_COUNT = 3
AXIS_ORDER = "[z, x, y]"
PLANE_XY = "xy"
PLANE_YZ = "yz"
PLANE_XZ = "xz"
@dataclass
class _OpenVolume:
"""Opened volume plus shape needed by the stitching tasks."""
path: Path
volume: VolumeBase
shape: tuple[int, int, int]
@dataclass(frozen=True)
class _PlaneSpec:
"""How a named orthogonal plane maps global axes to output axes."""
plane: str
assembly_axis: int
row_axis: int
col_axis: int
def _make_volume(volume_path: Path) -> VolumeBase:
"""Create the tomoscan volume reader matching the input path."""
suffix = volume_path.suffix.lower()
if suffix in (".tif", ".tiff"):
return MultiTIFFVolume(file_path=str(volume_path))
if volume_path.is_dir():
return TIFFVolume(folder=str(volume_path))
raise ValueError(
f"Unsupported volume format for {volume_path}. Expected TIFF file or directory."
)
def _read_volume_shape(volume_path: Path, volume: VolumeBase) -> tuple[int, int, int]:
"""Return the 3D shape of a volume."""
shape = tuple(int(axis_size) for axis_size in volume.get_volume_shape())
if len(shape) != 3:
raise RuntimeError(
f"Expected a 3D volume, got shape {shape} for {volume_path}."
)
return shape
@contextmanager
def _open_volume(volume_path: Path):
"""Open a volume and expose its shape."""
volume = _make_volume(volume_path)
shape = _read_volume_shape(volume_path, volume)
yield _OpenVolume(
path=volume_path,
volume=volume,
shape=shape,
)
def _load_plane_from_volume(volume: VolumeBase, plane: str, index: int) -> np.ndarray:
"""Load one XY, YZ, or XZ 2D plane as float32."""
if plane == PLANE_XY:
return np.asarray(volume.get_slice(axis=Z_AXIS, index=index), dtype=np.float32)
if plane == PLANE_YZ:
return np.asarray(volume.get_slice(axis=X_AXIS, index=index), dtype=np.float32)
if plane == PLANE_XZ:
return np.asarray(volume.get_slice(axis=Y_AXIS, index=index), dtype=np.float32)
raise ValueError(f"Unsupported plane: {plane!r}")
def _load_local_overlap_volumes(
reference_volume: _OpenVolume,
moving_volume: _OpenVolume,
center_reference_pixels: np.ndarray,
shift_pixels: np.ndarray,
window_shape: tuple[int, int, int],
) -> tuple[np.ndarray, np.ndarray]:
"""Load matching 3D windows around a guessed overlap location."""
ref_center = np.asarray(center_reference_pixels, dtype=int)
window_shape_array = np.asarray(window_shape, dtype=int)
def overlap_window(
length_ref: int, length_mov: int, shift: int, center: int, size: int
) -> tuple[int, int]:
ref_start, _, overlap = _compute_overlap_bounds(length_ref, length_mov, shift)
size = min(size, overlap)
start = max(ref_start, center - size // 2)
end = min(ref_start + overlap, start + size)
start = end - size
return start, end
z0, z1 = overlap_window(
reference_volume.shape[Z_AXIS],
moving_volume.shape[Z_AXIS],
int(shift_pixels[Z_AXIS]),
int(ref_center[Z_AXIS]),
int(window_shape_array[Z_AXIS]),
)
x0, x1 = overlap_window(
reference_volume.shape[X_AXIS],
moving_volume.shape[X_AXIS],
int(shift_pixels[X_AXIS]),
int(ref_center[X_AXIS]),
int(window_shape_array[X_AXIS]),
)
y0, y1 = overlap_window(
reference_volume.shape[Y_AXIS],
moving_volume.shape[Y_AXIS],
int(shift_pixels[Y_AXIS]),
int(ref_center[Y_AXIS]),
int(window_shape_array[Y_AXIS]),
)
reference_chunk = np.asarray(
reference_volume.volume.load_chunk(
(slice(z0, z1), slice(x0, x1), slice(y0, y1))
),
dtype=np.float32,
)
moving_chunk = np.asarray(
moving_volume.volume.load_chunk(
(
slice(z0 - shift_pixels[Z_AXIS], z1 - shift_pixels[Z_AXIS]),
slice(x0 - shift_pixels[X_AXIS], x1 - shift_pixels[X_AXIS]),
slice(y0 - shift_pixels[Y_AXIS], y1 - shift_pixels[Y_AXIS]),
)
),
dtype=np.float32,
)
return reference_chunk, moving_chunk
def _downsample_3d_xy(data: np.ndarray, factor: int) -> np.ndarray:
"""Average-pool the X/Y dimensions of a 3D block."""
if factor <= 1:
return data.astype(np.float32, copy=False)
rows = (data.shape[X_AXIS] // factor) * factor
cols = (data.shape[Y_AXIS] // factor) * factor
if rows == 0 or cols == 0:
return data.astype(np.float32, copy=False)
trimmed = data[:, :rows, :cols]
reshaped = trimmed.reshape(
data.shape[Z_AXIS], rows // factor, factor, cols // factor, factor
)
return reshaped.mean(axis=(2, 4), dtype=np.float32)
def _guess_shift_pixels(guess_mm: list[float], voxel_size_um: np.ndarray) -> np.ndarray:
"""Convert a [z, x, y] displacement from millimetres to pixels."""
guess_mm_array = np.asarray(guess_mm, dtype=np.float64)
if guess_mm_array.shape != (AXIS_COUNT,):
raise ValueError(f"guess must have shape {AXIS_ORDER}.")
if np.any(voxel_size_um <= 0):
raise ValueError(f"Invalid voxel_size_um: {voxel_size_um}")
return np.rint((guess_mm_array * 1000.0) / voxel_size_um).astype(int)
def _compute_overlap_bounds(
ref_length: int, mov_length: int, shift: int
) -> tuple[int, int, int]:
"""Return matching reference/moving starts and overlap length on one axis."""
global_start = max(0, shift)
global_end = min(ref_length, shift + mov_length)
overlap = global_end - global_start
if overlap <= 0:
raise ValueError(
f"No overlap for lengths ({ref_length}, {mov_length}) with shift {shift}."
)
ref_start = global_start
mov_start = global_start - shift
return ref_start, mov_start, overlap
def _normalize_array(data: np.ndarray) -> np.ndarray:
"""Center and scale an array before correlation."""
centered = data.astype(np.float32, copy=False) - float(data.mean())
std = float(centered.std())
if std < 1e-8:
return centered
return centered / std
def _correlate_volume_shift(
reference_volume: np.ndarray,
moving_volume: np.ndarray,
max_shift_pixels: np.ndarray,
) -> np.ndarray:
"""Return the best residual shift between two local 3D blocks."""
residual, _ = _correlate_volume_shift_with_score(
reference_volume, moving_volume, max_shift_pixels
)
return residual
def _correlate_volume_shift_with_score(
reference_volume: np.ndarray,
moving_volume: np.ndarray,
max_shift_pixels: np.ndarray,
) -> tuple[np.ndarray, float]:
"""Correlate two 3D blocks and return the best lag with its score."""
corr = correlate(
_normalize_array(reference_volume),
_normalize_array(moving_volume),
mode="full",
method="fft",
)
z_lags = correlation_lags(
reference_volume.shape[Z_AXIS], moving_volume.shape[Z_AXIS], mode="full"
)
x_lags = correlation_lags(
reference_volume.shape[X_AXIS], moving_volume.shape[X_AXIS], mode="full"
)
y_lags = correlation_lags(
reference_volume.shape[Y_AXIS], moving_volume.shape[Y_AXIS], mode="full"
)
mask = (
(np.abs(z_lags[:, None, None]) <= int(max_shift_pixels[Z_AXIS]))
& (np.abs(x_lags[None, :, None]) <= int(max_shift_pixels[X_AXIS]))
& (np.abs(y_lags[None, None, :]) <= int(max_shift_pixels[Y_AXIS]))
)
if not np.any(mask):
raise RuntimeError("Empty volume-correlation mask.")
masked_corr = np.where(mask, corr, -np.inf)
peak_index = np.unravel_index(int(np.argmax(masked_corr)), masked_corr.shape)
residual = np.asarray(
[
z_lags[peak_index[Z_AXIS]],
x_lags[peak_index[X_AXIS]],
y_lags[peak_index[Y_AXIS]],
],
dtype=int,
)
return residual, float(masked_corr[peak_index])
def _candidate_overlap_centers(
reference_shape: tuple[int, int, int],
moving_shape: tuple[int, int, int],
shift_pixels: np.ndarray,
) -> list[np.ndarray]:
"""Generate representative reference points inside the guessed overlap."""
bounds = [
_compute_overlap_bounds(
reference_shape[axis], moving_shape[axis], int(shift_pixels[axis])
)
for axis in range(AXIS_COUNT)
]
def position(axis: int, fraction: float) -> int:
ref_start, _, overlap = bounds[axis]
if overlap <= 1:
return ref_start
return ref_start + int(round((overlap - 1) * fraction))
fractions = [
(0.5, 0.5, 0.5),
(0.5, 0.25, 0.25),
(0.5, 0.25, 0.75),
(0.5, 0.75, 0.25),
(0.5, 0.75, 0.75),
(0.25, 0.5, 0.5),
(0.75, 0.5, 0.5),
(0.5, 0.5, 0.25),
(0.5, 0.5, 0.75),
(0.5, 0.25, 0.5),
(0.5, 0.75, 0.5),
]
centers = []
seen = set()
for fraction in fractions:
center = tuple(position(axis, fraction[axis]) for axis in range(AXIS_COUNT))
if center not in seen:
centers.append(np.asarray(center, dtype=int))
seen.add(center)
return centers
def _default_coarse_window_shape(search_radius_px: np.ndarray) -> tuple[int, int, int]:
"""Choose the coarse correlation block shape from the search radius."""
return (
max(24, 4 * int(search_radius_px[Z_AXIS]) + 9),
max(96, 4 * int(search_radius_px[X_AXIS]) + 33),
max(96, 4 * int(search_radius_px[Y_AXIS]) + 33),
)
def _coarse_search_radius(
search_radius_px: np.ndarray, coarse_factor: int
) -> np.ndarray:
"""Scale the X/Y search radius to the downsampled coarse blocks."""
return np.maximum(
1,
np.asarray(
[
int(search_radius_px[Z_AXIS]),
int(np.ceil(search_radius_px[X_AXIS] / coarse_factor)),
int(np.ceil(search_radius_px[Y_AXIS] / coarse_factor)),
]
),
)
def _coarse_refine_shift(
reference_volume: _OpenVolume,
moving_volume: _OpenVolume,
reference_shape: tuple[int, int, int],
moving_shape: tuple[int, int, int],
guess_shift_pixels: np.ndarray,
coarse_window_shape: tuple[int, int, int],
coarse_radius: np.ndarray,
coarse_factor: int,
) -> tuple[np.ndarray, np.ndarray]:
"""Find a coarse residual shift and the overlap center that produced it."""
best_score = -np.inf
best_coarse_residual = None
best_center = None
for candidate_center in _candidate_overlap_centers(
reference_shape, moving_shape, guess_shift_pixels
):
reference_chunk, moving_chunk = _load_local_overlap_volumes(
reference_volume=reference_volume,
moving_volume=moving_volume,
center_reference_pixels=candidate_center,
shift_pixels=guess_shift_pixels,
window_shape=coarse_window_shape,
)
coarse_reference_chunk = _downsample_3d_xy(reference_chunk, coarse_factor)
coarse_moving_chunk = _downsample_3d_xy(moving_chunk, coarse_factor)
coarse_residual, score = _correlate_volume_shift_with_score(
coarse_reference_chunk,
coarse_moving_chunk,
coarse_radius,
)
if score > best_score:
best_score = score
best_coarse_residual = coarse_residual
best_center = candidate_center
if best_coarse_residual is None or best_center is None:
raise RuntimeError("Could not refine shift from any overlap candidate window.")
return best_coarse_residual, best_center
def _fine_refine_shift(
reference_volume: _OpenVolume,
moving_volume: _OpenVolume,
center_reference_pixels: np.ndarray,
coarse_shift_pixels: np.ndarray,
search_radius_px: np.ndarray,
coarse_factor: int,
) -> np.ndarray:
"""Refine around the coarse shift using full-resolution local blocks."""
fine_window_shape = (
max(24, 4 * int(min(2, search_radius_px[Z_AXIS])) + 9),
max(64, 8 * coarse_factor + 33),
max(64, 8 * coarse_factor + 33),
)
reference_chunk, moving_chunk = _load_local_overlap_volumes(
reference_volume=reference_volume,
moving_volume=moving_volume,
center_reference_pixels=center_reference_pixels,
shift_pixels=coarse_shift_pixels,
window_shape=fine_window_shape,
)
fine_radius = np.asarray(
[
min(1, int(search_radius_px[Z_AXIS])),
max(1, coarse_factor),
max(1, coarse_factor),
],
dtype=int,
)
return _correlate_volume_shift(reference_chunk, moving_chunk, fine_radius)
def _refine_shift_pixels(
reference_volume: _OpenVolume,
moving_volume: _OpenVolume,
reference_shape: tuple[int, int, int],
moving_shape: tuple[int, int, int],
guess_shift_pixels: np.ndarray,
search_radius_px: np.ndarray,
downsample_factor: int,
correlation_window_shape_px: np.ndarray | None,
) -> np.ndarray:
"""Refine one pairwise volume shift from a guessed pixel displacement."""
guess_shift_pixels = np.asarray(guess_shift_pixels, dtype=int)
search_radius_px = np.asarray(search_radius_px, dtype=int)
coarse_factor = max(int(downsample_factor), 1)
if correlation_window_shape_px is None:
coarse_window_shape = _default_coarse_window_shape(search_radius_px)
else:
coarse_window_shape = (
int(correlation_window_shape_px[Z_AXIS]),
int(correlation_window_shape_px[X_AXIS]),
int(correlation_window_shape_px[Y_AXIS]),
)
coarse_radius = _coarse_search_radius(search_radius_px, coarse_factor)
coarse_residual, coarse_center = _coarse_refine_shift(
reference_volume=reference_volume,
moving_volume=moving_volume,
reference_shape=reference_shape,
moving_shape=moving_shape,
guess_shift_pixels=guess_shift_pixels,
coarse_window_shape=coarse_window_shape,
coarse_radius=coarse_radius,
coarse_factor=coarse_factor,
)
coarse_residual[X_AXIS:] *= coarse_factor
coarse_shift = guess_shift_pixels + coarse_residual
fine_residual = _fine_refine_shift(
reference_volume=reference_volume,
moving_volume=moving_volume,
center_reference_pixels=coarse_center,
coarse_shift_pixels=coarse_shift,
search_radius_px=search_radius_px,
coarse_factor=coarse_factor,
)
refined_shift = coarse_shift + fine_residual
logger.info(f"Refined shift [z, x, y] in pixels: {refined_shift.tolist()}")
return refined_shift
def _stitch_2d(
reference_plane: np.ndarray,
moving_plane: np.ndarray,
row_shift: int,
col_shift: int,
) -> np.ndarray:
"""Stitch two already-matched 2D planes."""
return _stitch_many_2d(
[reference_plane, moving_plane], [0, row_shift], [0, col_shift]
)
def _match_intensity_on_overlap(
canvas: np.ndarray,
canvas_weights: np.ndarray,
moving_plane: np.ndarray,
row0: int,
col0: int,
) -> np.ndarray:
"""Linearly match a moving plane to the existing canvas overlap."""
row1 = row0 + moving_plane.shape[0]
col1 = col0 + moving_plane.shape[1]
overlap = canvas_weights[row0:row1, col0:col1] > 0.0
if not np.any(overlap):
return moving_plane.astype(np.float32, copy=False)
reference = (
canvas[row0:row1, col0:col1][overlap]
/ canvas_weights[row0:row1, col0:col1][overlap]
)
moving = moving_plane[overlap].astype(np.float32, copy=False)
finite = np.isfinite(reference) & np.isfinite(moving)
if np.count_nonzero(finite) < 2:
return moving_plane.astype(np.float32, copy=False)
reference = reference[finite]
moving = moving[finite]
moving_std = float(moving.std())
if moving_std < 1e-6:
scale = 1.0
offset = float(reference.mean() - moving.mean())
else:
scale, offset = np.polyfit(moving, reference, deg=1)
return moving_plane.astype(np.float32, copy=False) * float(scale) + float(offset)
def _cosine_blend_weights(
overlap_mask: np.ndarray,
row_shift: int,
col_shift: int,
) -> np.ndarray:
"""Build smooth overlap weights for inserting a moving plane."""
if not np.any(overlap_mask):
return np.ones(overlap_mask.shape, dtype=np.float32)
rows = np.where(overlap_mask.any(axis=1))[0]
cols = np.where(overlap_mask.any(axis=0))[0]
row_weights = None
col_weights = None
if rows.size > 1 and row_shift != 0:
t = np.linspace(0.0, 1.0, rows.size, dtype=np.float32)
if row_shift < 0:
t = t[::-1]
row_weights = 0.5 - 0.5 * np.cos(np.pi * t)
if cols.size > 1 and col_shift != 0:
t = np.linspace(0.0, 1.0, cols.size, dtype=np.float32)
if col_shift < 0:
t = t[::-1]
col_weights = 0.5 - 0.5 * np.cos(np.pi * t)
weights = np.full(overlap_mask.shape, 0.5, dtype=np.float32)
if row_weights is not None and col_weights is not None:
weights[np.ix_(rows, cols)] = (
row_weights[:, None] + col_weights[None, :]
) / 2.0
elif row_weights is not None:
weights[rows, :] = row_weights[:, None]
elif col_weights is not None:
weights[:, cols] = col_weights[None, :]
weights[~overlap_mask] = 1.0
return weights
def _stitch_many_2d(
planes: list[np.ndarray], row_shifts: list[int], col_shifts: list[int]
) -> np.ndarray:
"""Stitch 2D slices extracted from volumes into one larger plane.
``planes`` are already-selected XY, YZ, or XZ slices. ``row_shifts`` and
``col_shifts`` place each slice in the output canvas in pixel units.
"""
if not planes:
raise ValueError("At least one plane is required for stitching.")
row_min = min(min(row_shifts), 0)
col_min = min(min(col_shifts), 0)
row_max = max(
row_shift + plane.shape[0] for plane, row_shift in zip(planes, row_shifts)
)
col_max = max(
col_shift + plane.shape[1] for plane, col_shift in zip(planes, col_shifts)
)
canvas = np.zeros((row_max - row_min, col_max - col_min), dtype=np.float32)
canvas_weights = np.zeros_like(canvas, dtype=np.float32)
for plane, row_shift, col_shift in zip(planes, row_shifts, col_shifts):
row0 = row_shift - row_min
col0 = col_shift - col_min
row1 = row0 + plane.shape[0]
col1 = col0 + plane.shape[1]
plane = _match_intensity_on_overlap(canvas, canvas_weights, plane, row0, col0)
target_weights = canvas_weights[row0:row1, col0:col1]
overlap = target_weights > 0.0
new_region = ~overlap
if np.any(overlap):
# Match intensity before blending so overlapping regions agree.
current = canvas[row0:row1, col0:col1] / np.maximum(target_weights, 1.0)
blend_weights = _cosine_blend_weights(
overlap, int(row_shift), int(col_shift)
)
blended = current * (1.0 - blend_weights) + plane * blend_weights
canvas[row0:row1, col0:col1][overlap] = blended[overlap]
canvas_weights[row0:row1, col0:col1][overlap] = 1.0
if np.any(new_region):
canvas[row0:row1, col0:col1][new_region] = plane[new_region]
canvas_weights[row0:row1, col0:col1][new_region] = 1.0
return canvas
_STITCHED_PLANE_SPECS = {
"stitched_xy": _PlaneSpec(
plane=PLANE_XY,
assembly_axis=Z_AXIS,
row_axis=X_AXIS,
col_axis=Y_AXIS,
),
"stitched_yz": _PlaneSpec(
plane=PLANE_YZ,
assembly_axis=X_AXIS,
row_axis=Z_AXIS,
col_axis=Y_AXIS,
),
"stitched_xz": _PlaneSpec(
plane=PLANE_XZ,
assembly_axis=Y_AXIS,
row_axis=Z_AXIS,
col_axis=X_AXIS,
),
}
def _stitch_matched_plane(
opened_volumes: list[_OpenVolume],
assembly_points_pixels: np.ndarray,
volume_shifts_pixels: np.ndarray,
spec: _PlaneSpec,
) -> np.ndarray:
"""Load matching slices from all volumes and stitch one output plane."""
return _stitch_many_2d(
[
_load_plane_from_volume(
opened_volume.volume,
spec.plane,
int(assembly_point[spec.assembly_axis]),
)
for opened_volume, assembly_point in zip(
opened_volumes, assembly_points_pixels
)
],
row_shifts=[int(shift[spec.row_axis]) for shift in volume_shifts_pixels],
col_shifts=[int(shift[spec.col_axis]) for shift in volume_shifts_pixels],
)
def _derive_output_path(volume_paths: list[Path]) -> Path:
"""Choose the default output directory from the input volume names."""
output_dir = volume_paths[0].parent.parent / "stitched"
if len(volume_paths) == 2:
stem = f"{volume_paths[0].stem}__{volume_paths[1].stem}_orthogonal_slices"
else:
stem = (
f"{volume_paths[0].stem}__to__{volume_paths[-1].stem}_"
f"{len(volume_paths)}vol_orthogonal_slices"
)
return output_dir / stem
def _save_stitched_slices(
output_path: Path,
stitched_slices: dict[str, np.ndarray],
volume_shifts_pixels: np.ndarray,
assembly_points_pixels: np.ndarray,
) -> None:
"""Write stitched TIFF planes and a small metadata JSON file."""
if output_path.suffix.lower() in (".tif", ".tiff"):
output_path = output_path.with_suffix("")
output_path.mkdir(parents=True, exist_ok=True)
metadata = {
"volume_shifts_pixels": volume_shifts_pixels.tolist(),
"assembly_points_pixels": assembly_points_pixels.tolist(),
"slices": {
"stitched_xy": "stitched_xy.tiff",
"stitched_yz": "stitched_yz.tiff",
"stitched_xz": "stitched_xz.tiff",
},
}
for plane_name, plane_data in stitched_slices.items():
imwrite(
output_path / metadata["slices"][plane_name],
plane_data.astype(np.float32, copy=False),
description=json.dumps({"name": plane_name}),
)
(output_path / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
def _normalize_volume_paths(volume_paths: list[Path] | None) -> list[Path]:
"""Validate and convert the ordered input volume paths."""
if volume_paths is None:
raise ValueError("Provide volume_paths.")
paths = [Path(path) for path in volume_paths]
if len(paths) < 2:
raise ValueError("At least two volume paths are required.")
return paths
def _normalize_xyz_rows(
value: object,
row_count: int,
*,
dtype: Any,
error_message: str,
allow_scalar: bool = False,
) -> np.ndarray:
"""Normalize scalar/vector/per-row inputs to an ``(N, 3)`` array."""
rows = np.asarray(value, dtype=dtype)
if allow_scalar and rows.shape == ():
rows = np.asarray([[rows.item()] * 3], dtype=dtype)
if rows.shape == (AXIS_COUNT,):
rows = np.repeat(rows[None, :], row_count, axis=0)
if rows.shape != (row_count, AXIS_COUNT):
raise ValueError(error_message)
return rows
def _normalize_guess_distances_mm(
guess_distances_mm: list[list[float]] | list[float] | None,
pair_count: int,
) -> np.ndarray:
"""Normalize pairwise guessed shifts in millimetres."""
if guess_distances_mm is None:
raise ValueError("Provide guess_distances_mm for the pairwise refinement.")
return _normalize_xyz_rows(
guess_distances_mm,
pair_count,
dtype=np.float64,
error_message=(
"Expected one [z, x, y] guess per consecutive pair, or one single "
"[z, x, y] guess to reuse for every pair."
),
)
def _normalize_search_radii_px(
search_radius_px: int | list[int] | list[list[int]], pair_count: int
) -> np.ndarray:
"""Normalize pairwise correlation search radii in pixels."""
return _normalize_xyz_rows(
search_radius_px,
pair_count,
dtype=int,
allow_scalar=True,
error_message=(
"search_radius_px must be an int, a [z, x, y] list, or one [z, x, y] "
"list per consecutive pair."
),
)
def _normalize_correlation_window_shapes_px(
correlation_window_shape_px: int | list[int] | list[list[int]] | None,
pair_count: int,
) -> list[np.ndarray | None]:
"""Normalize optional pairwise correlation window shapes in pixels."""
if correlation_window_shape_px is None:
return [None] * pair_count
shapes = _normalize_xyz_rows(
correlation_window_shape_px,
pair_count,
dtype=int,
allow_scalar=True,
error_message=(
"correlation_window_shape_px must be an int, a [z, x, y] list, "
"or one [z, x, y] list per consecutive pair."
),
)
if np.any(shapes <= 0):
raise ValueError("correlation_window_shape_px values must be positive.")
return [shape for shape in shapes]
def _normalize_voxel_sizes_um(
voxel_size_um: float | list[float] | list[list[float]],
volume_count: int,
) -> np.ndarray:
"""Normalize per-volume voxel sizes in micrometres."""
voxel_sizes = _normalize_xyz_rows(
voxel_size_um,
volume_count,
dtype=np.float64,
allow_scalar=True,
error_message=(
"voxel_size_um must be a scalar, one [z, x, y] list, or one "
"[z, x, y] list per volume."
),
)
if np.any(voxel_sizes <= 0):
raise ValueError(f"Invalid voxel_size_um: {voxel_sizes}")
return voxel_sizes
def _accumulate_pairwise_shifts(pair_shifts: np.ndarray) -> np.ndarray:
"""Convert consecutive-pair shifts to shifts relative to volume zero."""
volume_shifts = [np.zeros(AXIS_COUNT, dtype=int)]
for pair_shift in pair_shifts:
volume_shifts.append(volume_shifts[-1] + np.asarray(pair_shift, dtype=int))
return np.asarray(volume_shifts, dtype=int)
def _common_overlap_center_global_pixels(
volume_shapes: list[tuple[int, int, int]], volume_shifts_pixels: np.ndarray
) -> np.ndarray:
"""Find one global [z, x, y] point shared by all shifted volumes."""
center = []
for axis in range(AXIS_COUNT):
overlap_start = max(int(shift[axis]) for shift in volume_shifts_pixels)
overlap_end = min(
int(shift[axis]) + int(shape[axis])
for shift, shape in zip(volume_shifts_pixels, volume_shapes)
)
overlap = overlap_end - overlap_start
if overlap <= 0:
raise ValueError(f"No common overlap across all volumes on axis {axis}.")
center.append(overlap_start + overlap // 2)
return np.asarray(center, dtype=int)
[docs]
class RefineVolumePositionOutputsModel(BaseOutputModel):
volume_paths: list[str] = Field(
...,
description="Ordered list of volume paths used for refinement.",
)
pair_refined_shifts_pixels: list[list[int]] = Field(
...,
description="Refined consecutive-pair shifts in pixels ordered as [z, x, y].",
)
pair_refined_shifts_mm: list[list[float]] = Field(
...,
description=(
"Refined consecutive-pair shifts in millimetres ordered as [z, x, y]."
),
)
volume_shifts_pixels: list[list[int]] = Field(
...,
description=(
"Cumulative shifts of each volume relative to the first volume, ordered "
"as [z, x, y]. The first volume is always [0, 0, 0]."
),
)
volume_shifts_mm: list[list[float]] = Field(
...,
description=(
"Cumulative shifts of each volume relative to the first volume in "
"millimetres, ordered as [z, x, y]."
),
)
assembly_points_pixels: list[list[int]] = Field(
...,
description=(
"Local [z, x, y] assembly-point coordinates for each volume, all "
"corresponding to the same global overlap-center point."
),
)
[docs]
class RefineVolumePosition( # type: ignore[call-arg]
Task,
input_model=RefineVolumePositionInputsModel,
output_model=RefineVolumePositionOutputsModel,
):
"""Refine pairwise [z, x, y] shifts across an ordered list of volumes."""
[docs]
def run(self):
volume_paths = _normalize_volume_paths(
self.get_input_value("volume_paths", None)
)
pair_count = len(volume_paths) - 1
guess_distances_mm = _normalize_guess_distances_mm(
self.get_input_value("guess_distances_mm", None),
pair_count,
)
search_radii_px = _normalize_search_radii_px(
self.get_input_value("search_radius_px"),
pair_count,
)
correlation_window_shapes_px = _normalize_correlation_window_shapes_px(
self.get_input_value("correlation_window_shape_px", None),
pair_count,
)
voxel_sizes_um = _normalize_voxel_sizes_um(
self.get_input_value("voxel_size_um"),
len(volume_paths),
)
downsample_factor = int(self.get_input_value("downsample_factor"))
for path in volume_paths:
if not path.exists():
raise FileNotFoundError(f"Volume file not found: {path}")
with ExitStack() as stack:
opened_volumes = [
stack.enter_context(_open_volume(volume_path))
for volume_path in volume_paths
]
pair_refined_shifts_pixels = []
pair_refined_shifts_mm = []
for index in range(pair_count):
reference = opened_volumes[index]
moving = opened_volumes[index + 1]
pair_mean_voxel_size_um = 0.5 * (
voxel_sizes_um[index] + voxel_sizes_um[index + 1]
)
guessed_shift_pixels = _guess_shift_pixels(
guess_distances_mm[index].tolist(),
pair_mean_voxel_size_um,
)
refined_pair_shift_pixels = _refine_shift_pixels(
reference_volume=reference,
moving_volume=moving,
reference_shape=reference.shape,
moving_shape=moving.shape,
guess_shift_pixels=guessed_shift_pixels,
search_radius_px=search_radii_px[index],
downsample_factor=downsample_factor,
correlation_window_shape_px=correlation_window_shapes_px[index],
)
pair_refined_shifts_pixels.append(refined_pair_shift_pixels)
pair_refined_shifts_mm.append(
refined_pair_shift_pixels * pair_mean_voxel_size_um / 1000.0
)
pair_refined_shifts_pixels = np.asarray(
pair_refined_shifts_pixels, dtype=int
)
pair_refined_shifts_mm = np.asarray(
pair_refined_shifts_mm, dtype=np.float64
)
volume_shifts_pixels = _accumulate_pairwise_shifts(
pair_refined_shifts_pixels
)
volume_shifts_mm = np.vstack(
[
np.zeros(AXIS_COUNT, dtype=np.float64),
np.cumsum(pair_refined_shifts_mm, axis=0),
]
)
global_assembly_point = _common_overlap_center_global_pixels(
[volume.shape for volume in opened_volumes],
volume_shifts_pixels,
)
assembly_points_pixels = (
global_assembly_point[None, :] - volume_shifts_pixels
)
self.outputs.volume_paths = [str(path) for path in volume_paths]
self.outputs.pair_refined_shifts_pixels = pair_refined_shifts_pixels.tolist()
self.outputs.pair_refined_shifts_mm = pair_refined_shifts_mm.tolist()
self.outputs.volume_shifts_pixels = volume_shifts_pixels.tolist()
self.outputs.volume_shifts_mm = volume_shifts_mm.tolist()
self.outputs.assembly_points_pixels = assembly_points_pixels.tolist()
[docs]
class StitchMatchedSlicesOutputsModel(BaseOutputModel):
stitched_path: str = Field(
...,
description=(
"Output directory containing one stitched XY, one stitched YZ and one "
"stitched XZ TIFF slice."
),
)
[docs]
class StitchMatchedSlices( # type: ignore[call-arg]
Task,
input_model=StitchMatchedSlicesInputsModel,
output_model=StitchMatchedSlicesOutputsModel,
):
"""Assemble full stitched XY, YZ and XZ slices across all refined volumes."""
[docs]
def run(self):
volume_paths = _normalize_volume_paths(
self.get_input_value("volume_paths", None)
)
volume_shifts_pixels_input = self.get_input_value("volume_shifts_pixels", None)
assembly_points_pixels_input = self.get_input_value(
"assembly_points_pixels", None
)
if volume_shifts_pixels_input is None or assembly_points_pixels_input is None:
raise ValueError(
"Provide both volume_shifts_pixels and assembly_points_pixels."
)
volume_shifts_pixels = np.asarray(volume_shifts_pixels_input, dtype=int)
assembly_points_pixels = np.asarray(assembly_points_pixels_input, dtype=int)
output_path = self.get_input_value("output_path", None)
if volume_shifts_pixels.shape != (len(volume_paths), AXIS_COUNT):
raise ValueError(
"volume_shifts_pixels must have shape (n_volumes, 3) ordered as "
"[z, x, y]."
)
if assembly_points_pixels.shape != (len(volume_paths), AXIS_COUNT):
raise ValueError(
"assembly_points_pixels must have shape (n_volumes, 3) ordered as "
"[z, x, y]."
)
with ExitStack() as stack:
opened_volumes = [
stack.enter_context(_open_volume(volume_path))
for volume_path in volume_paths
]
stitched_slices = {
name: _stitch_matched_plane(
opened_volumes,
assembly_points_pixels,
volume_shifts_pixels,
spec,
)
for name, spec in _STITCHED_PLANE_SPECS.items()
}
if output_path is None:
output_path = _derive_output_path(volume_paths)
output_path = Path(output_path)
if output_path.suffix.lower() in (".tif", ".tiff"):
output_path = output_path.with_suffix("")
_save_stitched_slices(
output_path=output_path,
stitched_slices=stitched_slices,
volume_shifts_pixels=volume_shifts_pixels,
assembly_points_pixels=assembly_points_pixels,
)
stitched_shapes = {
name.removeprefix("stitched_"): data.shape
for name, data in stitched_slices.items()
}
logger.info(
"Saved orthogonal stitched slices to "
f"{output_path} with shapes {stitched_shapes}"
)
self.outputs.stitched_path = str(output_path)
__all__ = ["RefineVolumePosition", "StitchMatchedSlices"]
