Just in time computation of motion between two subsequent images using openpiv on a 2d + time Dask stack

I’m trying to compute the displacement of particules between two subsequent images and display the resulting displacement map as an overlay in Napari. I’m working with relatively large datasets and so preprocessing every frames in advance isn’t a very practical solution. The data is loaded as a 3d Dask array (2d + time) and I would like to lazily compute the displacement using the extended_search_area_piv function from the openpiv package in a similar fashion as demonstrated here. The openpiv function expects a pair of 2D frames as an argument. I attempted to use the map_blocks() from dask.array by rechunking the data along the time axis to single frames

stack = stack.rechunk({0:1})

and providing the piv function with two offset copies of the 2d+time array

area_search = partial(pyprocess.extended_search_area_piv, window_size=32, overlap=16, search_area_size=32)
piv_disp = da.map_blocks(area_search, 
              stack[1:],
              stack[:-1],
              chunks=(1, *x.shape), dtype='float64')

Once I try computing the result for a slice I get an error related to shape of the provided data.

~/.pyenv/versions/neuro_39/lib/python3.9/site-packages/OpenPIV-0.23.4-py3.9-macosx-11.2-x86_64.egg/openpiv/pyprocess.py in extended_search_area_piv()
    753 
    754     if (window_size > frame_a.shape[0]) or (window_size > frame_a.shape[1]):
--> 755         raise ValueError("window size cannot be larger than the image")
    756 
    757     # get field shape

ValueError: window size cannot be larger than the image

I’m assuming that this is a result of the 3rd dimension being added to the expected 2d frames. Is there a way to perform the computation on demand on 2d slices of the array?

Note that I already tried doing a stack of delayed function calls:

x, y = openpiv.pyprocess.get_coordinates(image_size=stack.shape[-2:], search_area_size=32, overlap=16)

@delayed
def displacement_map(frame_a, frame_b):
    u, v, s2n = pyprocess.extended_search_area_piv(frame_a, frame_b, window_size=32, overlap=16, search_area_size=32)
    displ = np.abs(np.dot(u, v))
    return displ

UV = da.stack([da.from_delayed(displacement_map(stack[tt].astype(np.int32),
                                                stack[tt+1].astype(np.int32)),
                               x.shape, dtype='float64')
               for tt in range(stack.shape[0]-1)], axis=0)

however the ‘precompile’ time in this approach turned out to be extremely long.

If you have any suggestions let me know Thanks!

3 Likes

not a complete answer… but this might be a good application for dask.map_overlap. The example in the docs there (mapping the derivative function to adjacent chunks of the array x) is quite similar to what you need to do here. As an example:

In [1]: import dask.array as da

In [2]: stack = da.zeros(shape=(10, 256, 256), chunks=(1,256,256))

In [3]: def show_shape(x):
            # just an example to show what dask is delivering
            # do your extended_search_area_piv processing here
   ...:     print(x.shape)
   ...:     return x
   ...:

# with depth=(1,0,0), dask will provide each timepoint, along with the prior & following timepoint
In [4]: out = stack.map_overlap(show_shape, depth=(1,0,0)).compute()
(0, 0, 0)
(1, 1, 1)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)
(3, 256, 256)

# but it then trims away the neighboring strips from the output to give you the final result
In [5]: out.shape
Out[5]: (10, 256, 256)

see the “boundary” argument for options on handling the edges.

4 Likes

Let us know how it goes when you give this a try, ok?

1 Like

It turns out the dask.map_overlap function does exactly what I’m looking for. Here is a snippet of the working solution:

x, y = openpiv.pyprocess.get_coordinates(image_size=stack[0].shape, search_area_size=32, overlap=16)

def piv_displacement(slices):
    u, v, s2n = openpiv.pyprocess.extended_search_area_piv(slices[1], slices[2], window_size=32, overlap=16, search_area_size=32)
    displ = np.abs(np.dot(u, v))[None]
    return displ

displacement_map = stack.astype(np.int32).map_overlap(piv_displacement, trim=False, depth={0: 1}, dtype=np.float64)

In this version I’m still taking 3 “slices” of the stack even though I’m only using two in the extended_search_area_piv function. From the documentation I’ve seen that it is possible to specify an asymmetric depth however I haven’t been able to get that to work for now. I will post an update when that happen.

2 Likes

Great. Glad it helped!

1 Like

Very cool, that’s great news!

1 Like