Add SBEM stage-position XY alignment option

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}')
 
 

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, 

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:

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',

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:

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),
+    }

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
+