[JAX] Import scipy submodules directly in code that uses them.

JAX will soon import scipy lazily, so code that relies on JAX having already imported submodules of scipy will fail. Change code to import what it uses.

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.