Fix and add more options for non-optical flow register_ROIs; add feature to interpolate shifts based on patch locations in tile_and_correct

[ethanbb]

Description

The problem

This addresses an issue I raised on the slack where the non-optical flow path of register_ROIs seems to have a mistake. It works by finding piecewise-rigid shifts between the two template images using tile_and_correct. Then, the matrix of shifts for each patch should be interpolated into a shift map of the same size as the image, to be used for cv2.remap. However, these lines:

CaImAn/caiman/base/rois.py

Lines 434 to 435 in bb55800

	x_remap = (-np.resize(shifts_x, dims) + x_grid).astype(np.float32)
	y_remap = (-np.resize(shifts_y, dims) + y_grid).astype(np.float32)

actually just repeat the entries of the matrix, in C-order, to fill in a new matrix of the requested shape, so the result will make no sense.

I used this test script: https://gist.github.com/ethanbb/12e6efa6ba31733f769c0b6c3f82ee98 to confirm that it gives weirdly warped results. Note that with the current code, setting use_opt_flow to False always gives an error, because the tile_and_correct call sets shifts_opencv=True, but this path does not return a value for xy_grid, which is then used on the next line. So, for testing, I temporarily changed shifts_opencv to False. Here are the results with the current code:

In this case, it seems like we should just stick with optical flow. However, one thing I've observed and which motivated this inquiry is that optical flow struggles with larger shifts. Here is what happens when I increase the amount of non-rigid offset from 2 to 6:

Clearly, neither method is satisfactory.

Solving the general problem

Up-sampling shifts from patches does not only happen in register_ROIs. It occurs in multiple places in the motion correction code as well, but for this PR I just focused on tile_and_correct and tile_and_correct_3d. In these functions, resizing occurs for both the OpenCV (to full image size in order to use cv2.remap) and non-OpenCV (to an Nx denser grid, depending on upsample_factor_grid) methods. I realized when looking into this that it was not being done in the most rigorously correct way; it is implemented as a standard image resize, but the patch centers at the borders do not coincide with the borders of the new image, and these centers may not even be evenly spaced (at the lower and right edges).

This imprecision may not usually be a big deal for large images, and may be necessary in order to correct whole movies efficiently. However, I thought that at least for the case of aligning just 2 images, I may as well add a way to up-sample more precisely by interpolating using the exact location of each patch. Thus, I added get_patch_centers and interpolate_shifts to interpolate shifts to/from specific pixel locations using scipy (I don't think OpenCV has an easy way of doing this). I also added the flag interp_shifts_precisely for the two tile_and_correct functions to enable this method. I haven't tested it on a full motion correction run yet, but it could potentially be promoted to a new parameter in the future.

Finally, I used the same functions to do the interpolation correctly within register_ROIs. Using the new get_patch_centers function, the xy_grid output is unnecessary, so it works with either setting of shifts_opencv. I also set interp_shifts_precisely to True here, and it seems to be working fine.

Results with nr_offset = 2:

And with nr_offset = 6:

Other changes to register_ROIs

• I added an optional argument align_options to allow overriding default motion correction params (for tile_and_correct method only). This is necessary, among other things, to increase max_deviation_rigid from 2 in order to deal with cases like the second example above.
• I added a clause to deal with the rigid case (if setting max_deviation_rigid = 0)
• For the optical flow branch, I noticed that the images are being multiplied by 255 and converted to uint8 twice in a row, causing lots of overflow. I assumed this was a bug, so I removed the second one.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)

Has your PR been tested?

Aside from the test script mentioned above, caimanmanager test passes. I'm trying to run demotest, having some issues with it but I'll update if I can get it working.

[pgunn]

Very nice work; do you mind cherry-picking the CI diffs from dev into this so tests can run?

[pgunn]

The "except" part probably could use some explanation

[ethanbb]

Very nice work; do you mind cherry-picking the CI diffs from dev into this so tests can run?

Thanks! Can you be more specific - this should already be applied on top of dev. Are the CI changes in another branch?

[pgunn]

Very nice work; do you mind cherry-picking the CI diffs from dev into this so tests can run?

Thanks! Can you be more specific - this should already be applied on top of dev. Are the CI changes in another branch?

They should be in dev now; the problem is github's test infrastructure runs on your branch, not on what the dev branch would look like with your changes, so your branch being slightly out of date means the fixes are not effective.

(rebasing would also work)

[ethanbb]

Sorry would you mind telling me the specific commits you're talking about? As far as I can tell it's not behind flatironinstitute/dev.

[pgunn]

Sorry would you mind telling me the specific commits you're talking about? As far as I can tell it's not behind flatironinstitute/dev.

You're right. Looks like it broke upstream again.

I'd like to merge this sooner and I'll have to look at the CI issue separately; I may need to reach out if it fails tests when they're working again.

[pgunn]

Are you ready for this to be merged now?

[ethanbb]

I think it's up to you - I have been using it fine on my own branch. I wasn't able to get the demo tests working locally. I think it might be worth considering whether that interp_shifts_precisely option should be made into a real parameter - I would like to try using it to motion correct a movie at some point and see how it differs in terms of runtime and results.

[pgunn]

Let me see if the CI stuff is quick to fix.

[pgunn]

I just updated dev with a CI fix.

[pgunn]

Passes CI now, which is good; you still sound uncertain on whether you want this style of parameterisation or something else; make a call on that. We could merge this now, or if you want to rework some of the approach first that's also fine.

[ethanbb]

OK, I tried motion correcting a long recording using interp_shifts_precisely. It made very little difference in the result, but it also didn't take much longer than normal, so I decided it's fine to add it as a new parameter (false by default, though).

[ethanbb]

Actually, before merging, I just want to update apply_shifts_movie to use the new option as well.

[ethanbb]

Not done yet, but when going through the code, I noticed a couple of things that seem to be bugs:

• When applying shifts to a 2D image in tile_and_correct using opencv remap, the border_nan option isn't taken into account (even though it is used when doing the same operation in apply_shifts_movie)
• In apply_shifts_movie, for the 3D case, the shifts aren't un-negated as they are for the 2D case (as far as I can tell there should be no difference)
• In apply_shifts_movie, the upsampling of the patch grid that happens when shifts_opencv is False is not taken into account (should cause an error when trying to reshape the shifts)

I'm going to assume these are actually bugs and fix them, but just a warning that if they're not I may be breaking something...

[ethanbb]

OK, now I'm done with this.

[ethanbb]

Should be all set now.