google-research/sofima

Scalable Optical Flow-based Image Montaging and Alignment

SOFIMA

SOFIMA (Scalable Optical Flow-based Image Montaging and Alignment) is a tool for stitching, aligning and warping large 2d, 3d and 4d microscopy datasets.

This is not an officially supported Google product.

Installation

SOFIMA is implemented purely in Python, and does not require a build step. To install it directly from the repository, run:

  pip install git+https://github.com/google-research/sofima

Overview

SOFIMA uses optical flow regularized with an elastic mesh to establish maps between data in different coordinate systems. Both the flow estimator as well as the mesh solver are implemented in JAX and will automatically take advantage of GPU acceleration if the hardware if available.

A core data structure used throughout the project is a coordinate map stored as a dense array of relative offsets (see the module docstring in map_utils.py for details). Among other uses, this is the representation of the estimated flow fields and the mesh node positions.

Example usage

• electron microscopy tile stitching
• electron microscopy section alignment

License

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this software except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

sofima's Issues

Aligning consecutive slices

Hello,

Use case

I have an input volume that is a reconstruction of 2D slices into a 3D volume. Therefore, the consecutive slices are jagged and the volume along perpendicular directions are not smooth. My idea is to try to align consecutive slices of that volume.

Code

To achieve that, I tried to use this package. Here is the snippet code I am using.

1. Could you tell me if I am using your repository right?
2. Do you have some advices concerning the different parameters (patch_size, step, stride, overlap, ...) along this pipeline?

from sofima.flow_field import JAXMaskedXCorrWithStatsCalculator
from sofima.flow_utils import clean_flow
from sofima.warp import ndimage_warp

patch_size = 15
step = 1
overlap = 3

maskcorr = JAXMaskedXCorrWithStatsCalculator()
flow = maskcorr.flow_field(img1, img2, patch_size, step)
clean_output = clean_flow(
    flow, 
    min_peak_ratio=np.min(flow[2]), 
    min_peak_sharpness=np.min(flow[1]), 
    max_magnitude=0, 
    max_deviation=0,
)
warped_img = ndimage_warp(
    img2, 
    clean_output, 
    stride=(step, step), 
    work_size=(patch_size, patch_size), 
    overlap=(overlap, overlap)
)

Results

I would like to apply this snippet of code iteratively. Concretely, this means:
a. Alignment between slice n and slice n+1 -> warping of slice n+1
b. Alignment between warped slice n+1 and slice n+2 -> warping of slice n+2
c. Alignment between warped slice n+2 and slice n+3 -> ...

However, currently, applying this iterative procedure is giving strange results: pixels are disappearing, see images below.

3. Do you have an idea why this is happening?

Thank you for your help.

Non-overlapping tiles

I am using a dataset that sometimes contains tiles that don't overlap. I've added an image of the coarse offsets. We can see that one offset (the offset between tiles (0, 0) and (0, 1) is not computed).

I use the workflow from the em_stitching Colab notebook. I compute the coarse offsets and mesh with this code:

from sofima import stitch_rigid
cx, cy = stitch_rigid.compute_coarse_offsets(grid_size, tile_map)

f, ax = plt.subplots(1, 2, figsize=(10, 5))
ax[0].quiver((0, 1, 2), (0, 1, 2), cx[0, 0, ...], cx[1, 0, ...])
ax[0].set_ylim(-0.5, 2.5)
ax[0].set_xlim(-0.5, 1.5)
ax[0].set_title('horizontal NNs')
ax[1].quiver((0, 1, 2), (0, 1, 2), cy[0, 0, ...], cy[1, 0, ...])
ax[1].set_ylim(-0.5, 1.5)
ax[1].set_xlim(-0.5, 2.5)
ax[1].set_title('vertical NNs')

coarse_mesh = stitch_rigid.optimize_coarse_mesh(cx, cy)

Then I use that to calculate the fine mesh with this code:

from sofima import stitch_elastic, flow_utils, mesh

stride = 20
cx = np.squeeze(cx)
cy = np.squeeze(cy)
fine_x, offsets_x = stitch_elastic.compute_flow_map(tile_map, cx, 0, stride=(stride, stride), batch_size=4)  # (x,y) -> (x+1,y)
fine_y, offsets_y = stitch_elastic.compute_flow_map(tile_map, cy, 1, stride=(stride, stride), batch_size=4)  # (x,y) -> (x,y+1)

# "min_peak_ratio": 1.4, "min_peak_sharpness": 1.4, "max_deviation": 5, "max_magnitude": 0}
kwargs = {"min_peak_ratio": 1.4, "min_peak_sharpness": 1.4, "max_deviation": 5, "max_magnitude": 0}
fine_x = {k: flow_utils.clean_flow(v[:, np.newaxis, ...], **kwargs)[:, 0, :, :] for k, v in fine_x.items()}
fine_y = {k: flow_utils.clean_flow(v[:, np.newaxis, ...], **kwargs)[:, 0, :, :] for k, v in fine_y.items()}

kwargs = {"min_patch_size": 10, "max_gradient": -1, "max_deviation": -1}
fine_x = {k: flow_utils.reconcile_flows([v[:, np.newaxis, ...]], **kwargs)[:, 0, :, :] for k, v in fine_x.items()}
fine_y = {k: flow_utils.reconcile_flows([v[:, np.newaxis, ...]], **kwargs)[:, 0, :, :] for k, v in fine_y.items()}

data_x = (cx, fine_x, offsets_x)
data_y = (cy, fine_y, offsets_y)

fx, fy, x, nbors, key_to_idx = stitch_elastic.aggregate_arrays(
    data_x, data_y, list(tile_map.keys()),
    coarse_mesh[:, 0, ...], stride=(stride, stride),
    tile_shape=next(iter(tile_map.values())).shape)

@jax.jit
def prev_fn(x):
    target_fn = ft.partial(stitch_elastic.compute_target_mesh, x=x, fx=fx,
                           fy=fy, stride=(stride, stride))
    x = jax.vmap(target_fn)(nbors)
    return jnp.transpose(x, [1, 0, 2, 3])

config = mesh.IntegrationConfig(dt=0.001, gamma=0., k0=0.01, k=0.1, stride=stride,
                                num_iters=1000, max_iters=20000, stop_v_max=0.001,
                                dt_max=100, prefer_orig_order=True,
                                start_cap=0.1, final_cap=10., remove_drift=True)

x, ekin, t = mesh.relax_mesh(x, None, config, prev_fn=prev_fn)

When I run the fine alignment cell, I get this error:

---

OverflowError Traceback (most recent call last)
in <cell line: 16>()
14
15 # Compute flow maps for horizontal and vertical directions
---> 16 fine_x, offsets_x = stitch_elastic.compute_flow_map(tile_map, cx, 0, stride=(stride, stride), batch_size=4) # (x,y) -> (x+1,y)
17 fine_y, offsets_y = stitch_elastic.compute_flow_map(tile_map, cy, 1, stride=(stride, stride), batch_size=4) # (x,y) -> (x,y+1)
18 /usr/local/lib/python3.10/dist-packages/sofima/stitch_elastic.py in compute_flow_map(tile_map, offset_map, axis, patch_size, stride, batch_size)
241 rounded_offset = stride[::-1] * np.round(offset / stride[::-1])
242
--> 243 overlap = -int(offset[axis])
244 overlap = pre.shape[1 - axis] - (
245 (pre.shape[1 - axis] - overlap) // stride[1 - axis] * stride[1 - axis]
OverflowError: cannot convert float infinity to integer

---

I believe the problem might have to do with non-overlapping tiles, because I have had the same problem earlier on a set where some of the tiles don't overlap, but with this specific grid I'm sure that this is not the case. When I check the tiles by hand in Fiji, they all seem to have overlapping features. Could someone help me to figure out how to solve this problem?

sofima.warp.render.tiles raise cv2.error: OpenCV(4.7.0) (-215:Assertion failed) dst.cols < SHRT_MAX && dst.rows < SHRT_MAX && src.cols < SHRT_MAX && src.rows < SHRT_MAX in function 'remap'

Following the stitching example for Sofima, the stitching is fine for small number of tiles, however I have 1631 EM tiles for stitching using the following tile_id_map

id_map = np.arange(1632).reshape(34, -1)
tile_id_map = id_map.copy() 
for i in range(id_map.shape[0]):
    if i % 2 == 1:
        tile_id_map[i, :] = tile_id_map[i, :][::-1]

Each tile is down sampled to about 5 MB

def read_image(image_path, x, y):
    f = open(image_path, 'rb')
    image_data = f.read()
    image = PIL.Image.open(io.BytesIO(image_data)).convert("L")
    image_np = np.array(image) 
    image.close()
    print(f"Read image {os.path.basename(image_path)}   \r", end=" ")
    return x, y, image_np 

with futures.ThreadPoolExecutor(max_workers=20) as executor:

    for y in range(tile_id_map.shape[0]):
        for x in range(tile_id_map.shape[1]):
            tile_id = tile_id_map[y, x]
                      
            image_path =  os.path.join(args.folder, f"down_sampled_Stitched_{tile_id:0>4}.tiff")
            f = executor.submit(read_image, image_path, x, y)
            future_to_tile_id.update({f: tile_id})


for f in futures.as_completed(future_to_tile_id):
    try:
        x, y, image_np = f.result()
        tile_map[(x, y)] = image_np
    except Exception as exc:
        print(f"Generated Execption: {exc}")

The rest of the stitching code is the same as Sofima stitching example using tile_space = (tile_id_map.shape[0],tile_id_map.shape[1]) and cx, cy = stitch_rigid.compute_coarse_offsets(tile_space, tile_map, ((1500, 500),(1500, 500)))

However, warp_subvolume raise error when trying to warp the tile using sofima.warp.render.tiles .
opencv version is 4.7

File "/home/youssef/workspace/sofima_trials/./stitchsupertiles.py", line 289, in <module>
    # Warp the tiles into a single image.
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/site-packages/sofima/warp.py", line 517, in render_tiles
    _render_tile(tile_x=x, tile_y=y, coord_map=coord_map)
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/site-packages/sofima/warp.py", line 465, in _render_tile
    warped_img, warped_mask = warp_subvolume(
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/site-packages/sofima/warp.py", line 182, in warp_subvolume
    f.result()
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/concurrent/futures/_base.py", line 451, in result
    return self.__get_result()
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/home/youssef/anaconda3/envs/sofima/lib/python3.10/site-packages/sofima/warp.py", line 169, in _warp_section
    warped[c, z, ...] = cvx2.remap(
cv2.error: OpenCV(4.7.0) /home/conda/feedstock_root/build_artifacts/libopencv_1678824110780/work/modules/imgproc/src/imgwarp.cpp:1724: error: (-215:Assertion failed) dst.cols < SHRT_MAX && dst.rows < SHRT_MAX && src.cols < SHRT_MAX && src.rows < SHRT_MAX in function 'remap'

There are several issues concerning opencv remap() dealing with large images. I am not sure if further down sampling of tiles would solve the issue and whether it is possible to pickle the result tile meshes for down sampled tiles and re-apply for large images.
Thanks in advance.