google-research/sofima

Scalable Optical Flow-based Image Montaging and Alignment

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.

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.

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?