from pathlib import Path
import numpy as np
from ewoks import execute_graph
from tifffile import TiffFile
from tomoscan.esrf.volume import MultiTIFFVolume
from ewokstomo.tasks.stitch_volumes import (
RefineVolumePosition,
StitchMatchedSlices,
_stitch_2d,
)
from ewokstomo.tests.test_utils import get_json_file
def _write_tiff_volume(
path: Path,
data: np.ndarray,
) -> Path:
volume = MultiTIFFVolume(
file_path=str(path),
data=data.astype(np.float32),
)
volume.save()
return path
def _read_stitched_tiff(
path: Path,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
if path.suffix.lower() in (".tif", ".tiff"):
path = path.with_suffix("")
with TiffFile(path / "stitched_xy.tiff") as tiff_file:
stitched_xy = tiff_file.asarray()
with TiffFile(path / "stitched_yz.tiff") as tiff_file:
stitched_yz = tiff_file.asarray()
with TiffFile(path / "stitched_xz.tiff") as tiff_file:
stitched_xz = tiff_file.asarray()
return stitched_xy, stitched_yz, stitched_xz
def _stitch_expected(
reference_plane: np.ndarray,
moving_plane: np.ndarray,
row_shift: int,
col_shift: int,
) -> np.ndarray:
row_min = min(0, row_shift)
col_min = min(0, col_shift)
row_max = max(reference_plane.shape[0], row_shift + moving_plane.shape[0])
col_max = max(reference_plane.shape[1], col_shift + moving_plane.shape[1])
canvas = np.zeros((row_max - row_min, col_max - col_min), dtype=np.float32)
weights = np.zeros_like(canvas, dtype=np.float32)
ref_r0 = -row_min
ref_c0 = -col_min
ref_r1 = ref_r0 + reference_plane.shape[0]
ref_c1 = ref_c0 + reference_plane.shape[1]
canvas[ref_r0:ref_r1, ref_c0:ref_c1] += reference_plane
weights[ref_r0:ref_r1, ref_c0:ref_c1] += 1.0
mov_r0 = row_shift - row_min
mov_c0 = col_shift - col_min
mov_r1 = mov_r0 + moving_plane.shape[0]
mov_c1 = mov_c0 + moving_plane.shape[1]
canvas[mov_r0:mov_r1, mov_c0:mov_c1] += moving_plane
weights[mov_r0:mov_r1, mov_c0:mov_c1] += 1.0
return canvas / np.maximum(weights, 1.0)
def _stitch_expected_many(
planes: list[np.ndarray], row_shifts: list[int], col_shifts: list[int]
) -> np.ndarray:
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)
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]
canvas[row0:row1, col0:col1] += plane
weights[row0:row1, col0:col1] += 1.0
return canvas / np.maximum(weights, 1.0)
[docs]
def test_stitch_2d_matches_intensity_before_blending():
reference = np.arange(36, dtype=np.float32).reshape(6, 6)
moving = reference[2:, :] * 2.0 + 10.0
stitched = _stitch_2d(reference, moving, row_shift=2, col_shift=0)
np.testing.assert_allclose(stitched[:6], reference, atol=1e-5)
[docs]
def test_refine_and_stitch_orthogonal_slices(tmp_path):
rng = np.random.default_rng(42)
scene = rng.standard_normal((14, 19, 21)).astype(np.float32) # (z, x, y)
true_shift_px = np.array([2, 3, -4], dtype=int)
reference = scene[:8, :10, 4:15]
moving = scene[
true_shift_px[0] : true_shift_px[0] + 8,
true_shift_px[1] : true_shift_px[1] + 10,
4 + true_shift_px[2] : 15 + true_shift_px[2],
]
volumes_dir = tmp_path / "volumes"
reference_path = _write_tiff_volume(
volumes_dir / "reference.tiff",
reference,
)
moving_path = _write_tiff_volume(
volumes_dir / "moving.tiff",
moving,
)
for guess_mm in ([0.20, 0.30, -0.40], [0.10, 0.10, -0.20], [0.30, 0.50, -0.60]):
refine_task = RefineVolumePosition(
inputs={
"volume_paths": [reference_path, moving_path],
"guess_distances_mm": guess_mm,
"search_radius_px": [2, 4, 4],
"downsample_factor": 1,
"voxel_size_um": [100.0, 100.0, 100.0],
}
)
refine_task.execute()
assert refine_task.outputs.volume_shifts_pixels == [[0, 0, 0], [2, 3, -4]]
assert refine_task.outputs.assembly_points_pixels == [[5, 6, 3], [3, 3, 7]]
output_path = tmp_path / "stitched"
stitch_task = StitchMatchedSlices(
inputs={
"volume_paths": [reference_path, moving_path],
"volume_shifts_pixels": refine_task.outputs.volume_shifts_pixels,
"assembly_points_pixels": refine_task.outputs.assembly_points_pixels,
"output_path": output_path,
}
)
stitch_task.execute()
stitched_xy, stitched_yz, stitched_xz = _read_stitched_tiff(output_path)
expected_xy = _stitch_expected(reference[5], moving[3], row_shift=3, col_shift=-4)
expected_yz = _stitch_expected(
reference[:, 6, :], moving[:, 3, :], row_shift=2, col_shift=-4
)
expected_xz = _stitch_expected(
reference[:, :, 3], moving[:, :, 7], row_shift=2, col_shift=3
)
np.testing.assert_allclose(stitched_xy, expected_xy, atol=1e-6)
np.testing.assert_allclose(stitched_yz, expected_yz, atol=1e-6)
np.testing.assert_allclose(stitched_xz, expected_xz, atol=1e-6)
[docs]
def test_stitch_volumes_workflow_forwards_volume_paths(tmp_path):
rng = np.random.default_rng(44)
scene = rng.standard_normal((14, 19, 21)).astype(np.float32) # (z, x, y)
true_shift_px = np.array([2, 3, -4], dtype=int)
reference = scene[:8, :10, 4:15]
moving = scene[
true_shift_px[0] : true_shift_px[0] + 8,
true_shift_px[1] : true_shift_px[1] + 10,
4 + true_shift_px[2] : 15 + true_shift_px[2],
]
volumes_dir = tmp_path / "volumes"
volume_paths = [
_write_tiff_volume(volumes_dir / "reference.tiff", reference),
_write_tiff_volume(volumes_dir / "moving.tiff", moving),
]
output_path = tmp_path / "workflow_stitched"
result = execute_graph(
str(get_json_file("stitch_volumes.json")),
inputs=[
{
"id": "refine_position_task",
"name": "volume_paths",
"value": volume_paths,
},
{
"id": "refine_position_task",
"name": "guess_distances_mm",
"value": [0.10, 0.10, -0.20],
},
{
"id": "refine_position_task",
"name": "search_radius_px",
"value": [2, 4, 4],
},
{
"id": "refine_position_task",
"name": "downsample_factor",
"value": 1,
},
{
"id": "refine_position_task",
"name": "voxel_size_um",
"value": [100.0, 100.0, 100.0],
},
{
"id": "stitch_task",
"name": "output_path",
"value": output_path,
},
],
)
assert result["stitched_path"] == str(output_path)
stitched_xy, stitched_yz, stitched_xz = _read_stitched_tiff(output_path)
expected_xy = _stitch_expected(reference[5], moving[3], row_shift=3, col_shift=-4)
expected_yz = _stitch_expected(
reference[:, 6, :], moving[:, 3, :], row_shift=2, col_shift=-4
)
expected_xz = _stitch_expected(
reference[:, :, 3], moving[:, :, 7], row_shift=2, col_shift=3
)
np.testing.assert_allclose(stitched_xy, expected_xy, atol=1e-6)
np.testing.assert_allclose(stitched_yz, expected_yz, atol=1e-6)
np.testing.assert_allclose(stitched_xz, expected_xz, atol=1e-6)
[docs]
def test_refine_and_stitch_tiff_volumes(tmp_path):
rng = np.random.default_rng(43)
scene = rng.standard_normal((14, 19, 21)).astype(np.float32) # (z, x, y)
true_shift_px = np.array([2, 3, -4], dtype=int)
reference = scene[:8, :10, 4:15]
moving = scene[
true_shift_px[0] : true_shift_px[0] + 8,
true_shift_px[1] : true_shift_px[1] + 10,
4 + true_shift_px[2] : 15 + true_shift_px[2],
]
volumes_dir = tmp_path / "volumes"
reference_path = _write_tiff_volume(volumes_dir / "reference.tiff", reference)
moving_path = _write_tiff_volume(volumes_dir / "moving.tiff", moving)
for path in (reference_path, moving_path):
path.with_name(f"{path.stem}_infos.txt").write_text("not parseable")
refine_task = RefineVolumePosition(
inputs={
"volume_paths": [reference_path, moving_path],
"guess_distances_mm": [0.10, 0.10, -0.20],
"search_radius_px": [2, 4, 4],
"downsample_factor": 1,
"voxel_size_um": [100.0, 100.0, 100.0],
}
)
refine_task.execute()
assert refine_task.outputs.volume_shifts_pixels == [[0, 0, 0], [2, 3, -4]]
assert refine_task.outputs.assembly_points_pixels == [[5, 6, 3], [3, 3, 7]]
stitch_task = StitchMatchedSlices(
inputs={
"volume_paths": [reference_path, moving_path],
"volume_shifts_pixels": refine_task.outputs.volume_shifts_pixels,
"assembly_points_pixels": refine_task.outputs.assembly_points_pixels,
}
)
stitch_task.execute()
output_path = Path(stitch_task.outputs.stitched_path)
assert output_path.is_dir()
stitched_xy, stitched_yz, stitched_xz = _read_stitched_tiff(output_path)
assert stitched_xy.shape == (13, 15)
assert stitched_yz.shape == (10, 15)
assert stitched_xz.shape == (10, 13)
expected_xy = _stitch_expected(reference[5], moving[3], row_shift=3, col_shift=-4)
expected_yz = _stitch_expected(
reference[:, 6, :], moving[:, 3, :], row_shift=2, col_shift=-4
)
expected_xz = _stitch_expected(
reference[:, :, 3], moving[:, :, 7], row_shift=2, col_shift=3
)
np.testing.assert_allclose(stitched_xy, expected_xy, atol=1e-6)
np.testing.assert_allclose(stitched_yz, expected_yz, atol=1e-6)
np.testing.assert_allclose(stitched_xz, expected_xz, atol=1e-6)
[docs]
def test_refine_and_stitch_orthogonal_slices_for_volume_list(tmp_path):
rng = np.random.default_rng(7)
scene = rng.standard_normal((20, 24, 30)).astype(np.float32) # (z, x, y)
volume0 = scene[:10, :12, 8:24]
volume1 = scene[2:12, 3:15, 4:20]
volume2 = scene[4:14, 6:18, :16]
volumes_dir = tmp_path / "volumes"
volume_paths = [
_write_tiff_volume(volumes_dir / "volume0.tiff", volume0),
_write_tiff_volume(volumes_dir / "volume1.tiff", volume1),
_write_tiff_volume(volumes_dir / "volume2.tiff", volume2),
]
refine_task = RefineVolumePosition(
inputs={
"volume_paths": volume_paths,
"guess_distances_mm": [[0.10, 0.20, -0.20], [0.30, 0.40, -0.50]],
"search_radius_px": [[2, 4, 4], [2, 4, 4]],
"downsample_factor": 1,
"voxel_size_um": [100.0, 100.0, 100.0],
}
)
refine_task.execute()
assert refine_task.outputs.pair_refined_shifts_pixels == [[2, 3, -4], [2, 3, -4]]
assert refine_task.outputs.volume_shifts_pixels == [
[0, 0, 0],
[2, 3, -4],
[4, 6, -8],
]
assert refine_task.outputs.assembly_points_pixels == [
[7, 9, 4],
[5, 6, 8],
[3, 3, 12],
]
output_path = tmp_path / "stitched_list"
stitch_task = StitchMatchedSlices(
inputs={
"volume_paths": volume_paths,
"volume_shifts_pixels": refine_task.outputs.volume_shifts_pixels,
"assembly_points_pixels": refine_task.outputs.assembly_points_pixels,
"output_path": output_path,
}
)
stitch_task.execute()
stitched_xy, stitched_yz, stitched_xz = _read_stitched_tiff(output_path)
expected_xy = _stitch_expected_many(
[volume0[7], volume1[5], volume2[3]],
row_shifts=[0, 3, 6],
col_shifts=[0, -4, -8],
)
expected_yz = _stitch_expected_many(
[volume0[:, 9, :], volume1[:, 6, :], volume2[:, 3, :]],
row_shifts=[0, 2, 4],
col_shifts=[0, -4, -8],
)
expected_xz = _stitch_expected_many(
[volume0[:, :, 4], volume1[:, :, 8], volume2[:, :, 12]],
row_shifts=[0, 2, 4],
col_shifts=[0, 3, 6],
)
np.testing.assert_allclose(stitched_xy, expected_xy, atol=1e-6)
np.testing.assert_allclose(stitched_yz, expected_yz, atol=1e-6)
np.testing.assert_allclose(stitched_xz, expected_xz, atol=1e-6)
