From bdfceee62638c42463ad1cacbbcf6b64df3721bd Mon Sep 17 00:00:00 2001 From: Marcel Sayre <52462209+msayr@users.noreply.github.com> Date: Tue, 30 Jun 2026 01:40:40 +1000 Subject: [PATCH] Add SBEM stage-position XY alignment option --- emalign/align_dataset_xy.py | 4 +++- emalign/align_xy/stitch_ongrid.py | 35 ++++++++++++++++++++++++++++++- emalign/io/sbem_image.py | 34 ++++++++++++++++++++++++++---- emalign/prep_config_xy.py | 12 ++++++++++- emalign/scripts/align_stack_xy.py | 18 ++++++++++++++-- tests/test_sbem_image.py | 15 +++++++++++++ tests/test_sbem_stage_offsets.py | 25 ++++++++++++++++++++++ 7 files changed, 134 insertions(+), 9 deletions(-) create mode 100644 tests/test_sbem_stage_offsets.py diff --git a/emalign/align_dataset_xy.py b/emalign/align_dataset_xy.py index 118a73a..fc35fd9 100644 --- a/emalign/align_dataset_xy.py +++ b/emalign/align_dataset_xy.py @@ -65,6 +65,7 @@ def align_dataset_xy(config_path, apply_clahe = main_config['apply_clahe'] stack_configs = main_config['stack_configs'] io_mode = main_config['io_mode'] + use_stage_positions = main_config.get('use_sbem_stage_positions', False) if not output_path.endswith('.zarr'): raise RuntimeError('Output path must be a zarr container (.zarr)') @@ -101,7 +102,8 @@ def align_dataset_xy(config_path, mongodb_config_filepath=mongodb_config_filepath, num_cores=num_workers, overwrite=overwrite, - wipe_progress_flag=wipe_this_stack) + wipe_progress_flag=wipe_this_stack, + use_stage_positions=use_stage_positions) logging.info(f'Done! Output can be found at: {output_path}') diff --git a/emalign/align_xy/stitch_ongrid.py b/emalign/align_xy/stitch_ongrid.py index 522f103..57fe5d9 100644 --- a/emalign/align_xy/stitch_ongrid.py +++ b/emalign/align_xy/stitch_ongrid.py @@ -13,7 +13,8 @@ def get_coarse_offset(tile_map, min_range=(10,100,0), min_overlap=5, filter_size=5, - overlap_cap=500): + overlap_cap=500, + tile_origins=None): ''' Compute coarse offset and mesh for initial rigid XY alignment ''' @@ -37,11 +38,43 @@ def get_coarse_offset(tile_map, cx[np.isinf(cx)] = np.nan cy[np.isinf(cy)] = np.nan + if tile_origins is not None: + cx, cy = _apply_tile_origin_offsets(cx, cy, tile_map, tile_origins) + coarse_mesh = stitch_rigid.optimize_coarse_mesh(cx, cy) return cx, cy, coarse_mesh +def _apply_tile_origin_offsets(cx, cy, tile_map, tile_origins): + """Replace neighbor offsets with metadata-derived tile origin deltas. + + SOFIMA's coarse mesh stores displacement relative to a no-overlap regular + grid. SBEM Image stage origins describe absolute tile origins, so adjacent + X links use ``dx - tile_width`` and adjacent Y links use ``dy - tile_height``. + """ + + tile_shape = next(iter(tile_map.values())).shape + height, width = tile_shape[:2] + cx = cx.copy() + cy = cy.copy() + + for (x, y), (origin_y, origin_x) in tile_origins.items(): + right = (x + 1, y) + if right in tile_origins: + right_y, right_x = tile_origins[right] + cx[0, 0, y, x] = right_x - origin_x - width + cx[1, 0, y, x] = right_y - origin_y + + down = (x, y + 1) + if down in tile_origins: + down_y, down_x = tile_origins[down] + cy[0, 0, y, x] = down_x - origin_x + cy[1, 0, y, x] = down_y - origin_y - height + + return cx, cy + + def get_elastic_mesh(tile_map, cx, cy, diff --git a/emalign/io/sbem_image.py b/emalign/io/sbem_image.py index f79caf4..0ef3472 100644 --- a/emalign/io/sbem_image.py +++ b/emalign/io/sbem_image.py @@ -12,6 +12,7 @@ _TILE_YX_POS = {} _TILE_YX_SOURCE = None _TILE_YX_PROJECT_ROOT = None +_TILE_STAGE_POS = {} def _clean_resolution(v): @@ -281,7 +282,7 @@ def _build_tile_yx_pos_map_from_imagelist(imagelist_path): entries.append((fname, tile_key, y, x)) if not entries: - return {} + return {}, {} y_vals = sorted({y for _, _, y, _ in entries}) y_to_row = {y: i for i, y in enumerate(y_vals)} @@ -307,19 +308,23 @@ def _build_tile_yx_pos_map_from_imagelist(imagelist_path): entry_cols = {fname: col - min_col for fname, col in entry_cols.items()} tile_map = {} + stage_map = {} for fname, tile_key, y, x in entries: pos = (y_to_row[y], entry_cols[fname]) + stage_pos = (y, x) tile_map[fname] = pos + stage_map[fname] = stage_pos if tile_key is not None: tile_map[tile_key] = pos + stage_map[tile_key] = stage_pos - return tile_map + return tile_map, stage_map def _ensure_tile_yx_pos_map(n): - global _TILE_YX_POS, _TILE_YX_SOURCE, _TILE_YX_PROJECT_ROOT + global _TILE_YX_POS, _TILE_YX_SOURCE, _TILE_YX_PROJECT_ROOT, _TILE_STAGE_POS lookup_keys = _lookup_keys(n) project_root = _find_project_root(n) @@ -344,10 +349,11 @@ def _ensure_tile_yx_pos_map(n): raise FileNotFoundError(f"No imagelist_*.txt files found in {logs_dir}") for imagelist_path in imagelist_files: - tile_map = _build_tile_yx_pos_map_from_imagelist(imagelist_path) + tile_map, stage_map = _build_tile_yx_pos_map_from_imagelist(imagelist_path) if any(key in tile_map for key in lookup_keys): _TILE_YX_POS = tile_map + _TILE_STAGE_POS = stage_map _TILE_YX_SOURCE = imagelist_path _TILE_YX_PROJECT_ROOT = project_root return @@ -376,6 +382,26 @@ def parse_yx_pos_from_name(n): ) +def get_stage_origin_from_name(n): + """Return raw SBEM Image stage origin as (y, x) for a tile image.""" + + _ensure_tile_yx_pos_map(n) + + for key in _lookup_keys(n): + if key in _TILE_STAGE_POS: + return _TILE_STAGE_POS[key] + + raise KeyError( + f"No stage origin found for {_basename(n)!r} in {_TILE_YX_SOURCE}" + ) + + +def get_stage_origins(tile_map_paths): + """Return stage origins keyed by tile position for an SBEM Image tile map.""" + + return {tile_pos: get_stage_origin_from_name(path) for tile_pos, path in tile_map_paths.items()} + + def parse_slice_from_name(n): """ Extract z-slice index from filenames like: diff --git a/emalign/prep_config_xy.py b/emalign/prep_config_xy.py index ad680dc..fdb7234 100644 --- a/emalign/prep_config_xy.py +++ b/emalign/prep_config_xy.py @@ -45,6 +45,7 @@ def prep_align_stacks(main_dir, num_workers, port, io_mode='volumescope', + use_sbem_stage_positions=False, project_name=None, force_overwrite=False): from emalign.align_xy.prep import ( @@ -54,6 +55,9 @@ def prep_align_stacks(main_dir, ) from emalign.io.backend import get_io_backend + if use_sbem_stage_positions and io_mode != 'sbem_image': + raise ValueError('--use-sbem-stage-positions can only be used with --mode sbem_image') + io_backend = get_io_backend(io_mode) # Make config dir @@ -175,7 +179,8 @@ def prep_align_stacks(main_dir, 'stride': stride, 'apply_gaussian': apply_gaussian, 'apply_clahe': apply_clahe, - 'io_mode': io_mode + 'io_mode': io_mode, + 'use_sbem_stage_positions': use_sbem_stage_positions } with open(os.path.join(config_dir, 'main_config.json'), 'w') as f: @@ -283,6 +288,11 @@ def main(): action='store_true', default=False, help='Force overwrite of existing config files. Default: user is prompted if configs exist') + parser.add_argument('--use-sbem-stage-positions', + dest='use_sbem_stage_positions', + action='store_true', + default=False, + help='Use SBEM Image x/y stage coordinates as the coarse XY tile-position prior. Only valid with --mode sbem_image.') parser.add_argument('--mode', metavar='MODE', dest='io_mode', diff --git a/emalign/scripts/align_stack_xy.py b/emalign/scripts/align_stack_xy.py index 0c11996..7d09dc1 100644 --- a/emalign/scripts/align_stack_xy.py +++ b/emalign/scripts/align_stack_xy.py @@ -49,7 +49,8 @@ def align_stack_xy(output_path, mongodb_config_filepath=None, num_cores=1, overwrite=False, - wipe_progress_flag=False): + wipe_progress_flag=False, + use_stage_positions=False): '''Align and stitch image stack in XY. @@ -73,6 +74,8 @@ def align_stack_xy(output_path, client = get_mongo_client(mongodb_config_filepath) db = get_mongo_db(client, project_name) io_backend = get_io_backend(io_mode) + if use_stage_positions and io_mode != 'sbem_image': + raise ValueError('Stage-position assisted XY alignment is only available in sbem_image mode') if wipe_progress_flag: logging.info(f"Wiping progress for stack: {stack_name}") @@ -154,6 +157,16 @@ def align_stack_xy(output_path, pbar.set_description(f'{stack.stack_name}: Loading tile_map...') tm = stack.get_tile_map(z, apply_gaussian, apply_clahe) tile_map = tm.tile_map + tile_origins = None + if use_stage_positions: + stage_origins = io_backend.get_stage_origins(tm.tile_map_paths) + y_res, x_res = resolution + min_y = min(y for y, _ in stage_origins.values()) + min_x = min(x for _, x in stage_origins.values()) + tile_origins = { + tile_pos: ((stage_y - min_y) / y_res, (stage_x - min_x) / x_res) + for tile_pos, (stage_y, stage_x) in stage_origins.items() + } metadata = {} if len(tile_map) > 1: @@ -176,7 +189,8 @@ def align_stack_xy(output_path, tm.tile_space, overlap=overlap, overlap_pad=80, - overlap_cap=1000 + overlap_cap=1000, + tile_origins=tile_origins ) if overlap > 160: diff --git a/tests/test_sbem_image.py b/tests/test_sbem_image.py index 96c70d5..2d85d63 100644 --- a/tests/test_sbem_image.py +++ b/tests/test_sbem_image.py @@ -36,6 +36,7 @@ def _reset_cache(): sbem_image._TILE_YX_POS = {} sbem_image._TILE_YX_SOURCE = None sbem_image._TILE_YX_PROJECT_ROOT = None + sbem_image._TILE_STAGE_POS = {} def test_backend_registration_exposes_sbem_image(): @@ -151,3 +152,17 @@ def test_sbem_stack_name_includes_grid_and_tile(tmp_path): sbem_image.get_stack_name(project / "tiles" / "g0001" / "t0000") == "g0001_t0000" ) + + +def test_stage_origins_are_available_by_tile_position(tmp_path): + _reset_cache() + project = _write_project(tmp_path) + tile_map_paths = { + (0, 0): project / "tiles" / "g0000" / "t0000" / "sample_g0000_t0000_s00001.tif", + (1, 0): project / "tiles" / "g0000" / "t0001" / "sample_g0000_t0001_s00001.tif", + } + + assert sbem_image.get_stage_origins(tile_map_paths) == { + (0, 0): (-619022, -608632), + (1, 0): (-619022, -421332), + } diff --git a/tests/test_sbem_stage_offsets.py b/tests/test_sbem_stage_offsets.py new file mode 100644 index 0000000..bb570cb --- /dev/null +++ b/tests/test_sbem_stage_offsets.py @@ -0,0 +1,25 @@ +import numpy as np +from emalign.align_xy.stitch_ongrid import _apply_tile_origin_offsets + + +def test_apply_tile_origin_offsets_converts_origins_to_sofima_offsets(): + tile_map = { + (0, 0): np.zeros((100, 200), dtype=np.uint8), + (1, 0): np.zeros((100, 200), dtype=np.uint8), + (0, 1): np.zeros((100, 200), dtype=np.uint8), + } + cx = np.zeros((2, 1, 2, 2), dtype=float) + cy = np.zeros((2, 1, 2, 2), dtype=float) + tile_origins = { + (0, 0): (0, 0), + (1, 0): (3, 180), + (0, 1): (90, -4), + } + + cx, cy = _apply_tile_origin_offsets(cx, cy, tile_map, tile_origins) + + assert cx[0, 0, 0, 0] == -20 + assert cx[1, 0, 0, 0] == 3 + assert cy[0, 0, 0, 0] == -4 + assert cy[1, 0, 0, 0] == -10 +