Adding support for Zarr stitching/fusion

coarse_registration.py

@@ -0,0 +1,120 @@
+import functools as ft
+import gc
+import jax
+import numpy as np
+import tensorstore as ts
+
+from sofima import flow_field
+
+
+QUERY_R_ORTHO = 100
+QUERY_OVERLAP_OFFSET = 0  # Overlap = 'starting line' in neighboring tile
+QUERY_R_OVERLAP = 100
+
+SEARCH_OVERLAP = 300  # Boundary - overlap = 'starting line' in search tile
+SEARCH_R_ORTHO = 100
+
+
+@ft.partial(jax.jit)
+def _estimate_relative_offset_zyx(base, 
+                                  kernel
+                                  ) -> list[float, float, float]:
+    # Calculate FFT: left = base, right = kernel
+    xc = flow_field.masked_xcorr(base, kernel, use_jax=True, dim=3)
+    xc = xc.astype(np.float32)
+    xc = xc[None, ...]
+
+    # Find strongest peak in FFT, pass in FFT image center
+    r = flow_field._batched_peaks(xc, 
+                                  ((xc.shape[1] + 1) // 2, (xc.shape[2] + 1) // 2, xc.shape[3] // 2), 
+                                  min_distance=2, 
+                                  threshold_rel=0.5)
+
+    # r returns a list, relative offset is here
+    relative_offset_xyz = r[0][0:3]
+    return [relative_offset_xyz[2], relative_offset_xyz[1], relative_offset_xyz[0]]
+
+
+def _estimate_h_offset_zyx(left_tile: ts.TensorStore, 
+                           right_tile: ts.TensorStore
+                           ) -> tuple[list[float], float]:
+    tile_size_xyz = left_tile.shape
+    mz = tile_size_xyz[2] // 2
+    my = tile_size_xyz[1] // 2
+
+    # Search Space, fixed
+    left = left_tile[tile_size_xyz[0]-SEARCH_OVERLAP:,
+                     my-SEARCH_R_ORTHO:my+SEARCH_R_ORTHO,
+                     mz-SEARCH_R_ORTHO:mz+SEARCH_R_ORTHO].read().result().T
+
+    # Query Patch, scanned against search space
+    right = right_tile[QUERY_OVERLAP_OFFSET:QUERY_OVERLAP_OFFSET + QUERY_R_OVERLAP*2,
+                       my-QUERY_R_ORTHO:my+QUERY_R_ORTHO,
+                       mz-QUERY_R_ORTHO:mz+QUERY_R_ORTHO].read().result().T
+
+    start_zyx = np.array(left.shape) // 2 - np.array(right.shape) // 2
+    pc_init_zyx = np.array([0, 0, tile_size_xyz[0] - SEARCH_OVERLAP + start_zyx[2]])
+    pc_zyx = np.array(_estimate_relative_offset_zyx(left, right))
+
+    return pc_init_zyx + pc_zyx
+
+
+def _estimate_v_offset_zyx(top_tile: ts.TensorStore, 
+                           bot_tile: ts.TensorStore,
+                          ) -> tuple[list[float], float]:
+    tile_size_xyz = top_tile.shape
+    mz = tile_size_xyz[2] // 2
+    mx = tile_size_xyz[0] // 2
+
+    top = top_tile[mx-SEARCH_R_ORTHO:mx+SEARCH_R_ORTHO, 
+                   tile_size_xyz[1]-SEARCH_OVERLAP:, 
+                   mz-SEARCH_R_ORTHO:mz+SEARCH_R_ORTHO].read().result().T  
+    bot = bot_tile[mx-QUERY_R_ORTHO:mx+QUERY_R_ORTHO, 
+                   0:QUERY_R_OVERLAP*2, 
+                   mz-QUERY_R_ORTHO:mz+QUERY_R_ORTHO].read().result().T
+
+    start_zyx = np.array(top.shape) // 2 - np.array(bot.shape) // 2
+    pc_init_zyx = np.array([0, tile_size_xyz[1] - SEARCH_OVERLAP + start_zyx[1], 0])    
+    pc_zyx = np.array(_estimate_relative_offset_zyx(top, bot))
+
+    return pc_init_zyx + pc_zyx
+
+
+def compute_coarse_offsets(tile_layout: np.ndarray, 
+                           tile_volumes: list[ts.TensorStore]
+                           ) -> tuple[np.ndarray, np.ndarray]:
+    layout_y, layout_x = tile_layout.shape
+
+    # Output Containers, sofima uses cartesian convention
+    conn_x = np.full((3, 1, layout_y, layout_x), np.nan)
+    conn_y = np.full((3, 1, layout_y, layout_x), np.nan)
+
+    # Row Pairs
+    for y in range(layout_y): 
+        for x in range(layout_x - 1):  # Stop one before the end 
+            left_id = tile_layout[y, x]
+            right_id = tile_layout[y, x + 1]
+            left_tile = tile_volumes[left_id]
+            right_tile = tile_volumes[right_id]
+
+            conn_x[:, 0, y, x] = _estimate_h_offset_zyx(left_tile, right_tile)
+            gc.collect()
+
+            print(f'Left Id: {left_id}, Right Id: {right_id}')
+            print(f'Left: ({y}, {x}), Right: ({y}, {x + 1})', conn_x[:, 0, y, x])
+
+    # Column Pairs -- Reversed Loops
+    for x in range(layout_x):
+        for y in range(layout_y - 1):
+            top_id = tile_layout[y, x]
+            bot_id = tile_layout[y + 1, x]
+            top_tile = tile_volumes[top_id]
+            bot_tile = tile_volumes[bot_id]
+
+            conn_y[:, 0, y, x] = _estimate_v_offset_zyx(top_tile, bot_tile)
+            gc.collect()
+
+            print(f'Top Id: {top_id}, Bottom Id: {bot_id}')
+            print(f'Top: ({y}, {x}), Bot: ({y + 1}, {x})', conn_y[:, 0, y, x])
+
+    return conn_x, conn_y

processor/warp.py

@@ -36,36 +36,29 @@ class StitchAndRender3dTiles(subvolume_processor.SubvolumeProcessor):
   """Renders a volume by stitching 3d tiles placed on a 2d grid."""
 
   _tile_meshes = None
-  _tile_idx_to_xy = None
+  _mesh_index_to_xy = {}
   _tile_boxes = {}
   _inverted_meshes = {}
 
   crop_at_borders = False
 
   def __init__(
       self,
-      *,
-      tile_map: Sequence[Sequence[int]],
-      tile_mesh_path: str,
-      tile_pattern_path: str,
+      tile_layout: Sequence[Sequence[int]],
+      tile_mesh: str,
+      xy_to_mesh_index: dict[int, tuple],
       stride: ZYX,
       offset: XYZ = (0, 0, 0),
       margin: int = 0,
       work_size: XYZ = (128, 128, 128),
       order: int = 1,
-      parallelism: int = 16,
-      input_volinfo=None,
+      parallelism: int = 16
   ):
     """Constructor.
 
     Args:
       tile_map: yx-shaped grid of tile IDs
-      tile_mesh_path: path to a npz file containing 'key_to_idx' and 'x' arrays,
-        as generated by `stitch_elastic.aggregate_arrays` and `mesh.solve_mesh`,
-        respectively
-      tile_pattern_path: volinfo path for the volumes containing individual
-        tiles; must contain '{tile_id}', which will be substituted with values
-        from `tile_map`
+      tile_idx_to_xy: index 
       stride: ZYX stride of the mesh in pixels
       offset: XYZ global offset to apply to the rendered image
       margin: number of pixels away from the tile boundary to ignore during
@@ -75,23 +68,27 @@ def __init__(
       work_size: see `warp.ndimage_warp`
       order: see `warp.ndimage_warp`
       parallelism: see `warp.ndimage_warp`
-      input_volinfo: not used
     """
-    del input_volinfo
-    self._tile_map = np.array(tile_map)
-    self._tile_mesh_path = tile_mesh_path
-    self._tile_pattern_path = tile_pattern_path
+    self._tile_layout = tile_layout
     self._stride = stride
     self._offset = offset
     self._margin = margin
     self._order = order
     self._parallelism = parallelism
     self._work_size = work_size
 
-    self._key_to_idx = {}
-    for y, row in enumerate(tile_map):
+    StitchAndRender3dTiles._tile_meshes = tile_mesh
+    StitchAndRender3dTiles._mesh_index_to_xy = {
+      v:k for k, v in xy_to_mesh_index.items()
+    }
+    assert StitchAndRender3dTiles._tile_meshes.shape[1] == len(
+        StitchAndRender3dTiles._mesh_index_to_xy
+    )
+
+    self._xy_to_tile_id = {}
+    for y, row in enumerate(tile_layout):
       for x, tile_id in enumerate(row):
-        self._key_to_idx[(x, y)] = tile_id
+        self._xy_to_tile_id[(x, y)] = tile_id
 
   def _open_tile_volume(self, tile_id: int) -> Any:
     """Returns a ZYX-shaped ndarray-like object representing the tile data."""
@@ -109,7 +106,7 @@ def _collect_tile_boxes(self, tile_shape_zyx: ZYX):
     )
 
     for i in range(StitchAndRender3dTiles._tile_meshes.shape[1]):
-      tx, ty = StitchAndRender3dTiles._tile_idx_to_xy[i]
+      tx, ty = StitchAndRender3dTiles._mesh_index_to_xy[i]
 
       mesh = StitchAndRender3dTiles._tile_meshes[:, i, ...]
       tg_box = map_utils.outer_box(mesh, map_box, self._stride)
@@ -141,9 +138,9 @@ def _get_dts(self, shape: ZYX, tx: int, ty: int) -> np.ndarray:
     mask = np.zeros(shape[1:], dtype=bool)
     if self._margin > 0:
       x0 = self._margin if tx > 0 else 0
-      x1 = -self._margin if tx < self._tile_map.shape[-1] - 1 else -1
+      x1 = -self._margin if tx < self._tile_layout.shape[-1] - 1 else -1
       y0 = self._margin if ty > 0 else 0
-      y1 = -self._margin if ty < self._tile_map.shape[-2] - 1 else -1
+      y1 = -self._margin if ty < self._tile_layout.shape[-2] - 1 else -1
       mask[y0:y1, x0:x1] = 1
     else:
       mask[...] = 1
@@ -178,7 +175,7 @@ def _load_tile_images(
       logging.info('Processing source %r (%r)', i, out_box)
 
       coord_map = StitchAndRender3dTiles._tile_meshes[:, i, ...]
-      tx, ty = StitchAndRender3dTiles._tile_idx_to_xy[i]
+      tx, ty = StitchAndRender3dTiles._mesh_index_to_xy[i]
 
       if i not in StitchAndRender3dTiles._inverted_meshes:
         # Add context to avoid rounding issues in map inversion.
@@ -209,8 +206,8 @@ def _load_tile_images(
       local_rel_box = sub_box.translate(-out_box.start)
       local_warp_box = local_rel_box.translate(local_out_box.start)
 
-      # Part of the inverted mesh that is needed to render the current
-      # region of interest.
+      # Part of the inverted mesh that is needed to render 
+      # the current region of interest.
       s = 1.0 / np.array(self._stride)[::-1]
       local_map_box = local_warp_box.scale(s).adjusted_by(
           start=(-2, -2, -2), end=(2, 2, 2)
@@ -251,30 +248,14 @@ def process(
     box = subvol.bbox
     logging.info('Processing %r', box)
 
-    mesh_init = False
-
-    if StitchAndRender3dTiles._tile_meshes is None:
-      data_path = self._tile_mesh_path
-      with file.Open(data_path, 'rb') as f:
-        data = np.load(f, allow_pickle=True)
-        StitchAndRender3dTiles._tile_idx_to_xy = {
-            v: k for k, v in data['key_to_idx'].item().items()
-        }
-        StitchAndRender3dTiles._tile_meshes = data['x']
-        assert StitchAndRender3dTiles._tile_meshes.shape[1] == len(
-            StitchAndRender3dTiles._tile_idx_to_xy
-        )
-      mesh_init = True
-
     volstores = {}
     for i in range(StitchAndRender3dTiles._tile_meshes.shape[1]):
-      tile_id = self._key_to_idx[StitchAndRender3dTiles._tile_idx_to_xy[i]]
+      tile_id = self._xy_to_tile_id[StitchAndRender3dTiles._mesh_index_to_xy[i]]
       volstores[i] = self._open_tile_volume(tile_id)
 
     # Bounding boxes representing a single tile placed the origin.
     tile_shape_zyx = next(iter(volstores.values())).shape
-    if mesh_init:
-      self._collect_tile_boxes(tile_shape_zyx)
+    self._collect_tile_boxes(tile_shape_zyx)
 
     # For blending, accumulate (weighted) image data as floats. This will
     # be normalized and cast to the desired output type once the image is

setup.cfg

@@ -28,6 +28,7 @@ install_requires =
   opencv-python>=4.5.5.62
   scipy>=1.2.3
   scikit-image>=0.17.2
+  tensorstore>=0.1.39
 
 [options.packages.find]
 where = .

stitch_elastic.py

@@ -134,8 +134,8 @@ def compute_flow_map3d(
       curr_box = bounding_box.BoundingBox(start=(0, 0, 0), size=tile_shape)
       nbor_box = bounding_box.BoundingBox(
           start=(
-              tile_shape[0] * (1 - axis) + offset[0],
-              tile_shape[1] * axis + offset[1],
+              offset[0],
+              offset[1],
               offset[2],
           ),
           size=tile_shape,