[BUG] Occasional glitchy mask during dynamics

[mmuunniiaakk]

I first posted this on the ImageJ forum a couple days ago, but realized this is the more appropriate place.

I have a managed to tweak a CP-SAM model (trained w/ v4.1.1) that performs very nicely on my 3D datasets. However, I recently noticed a bizarre glitch where occasional masks within the volume are incompletely filled--despite probs/flows looking good--resulting in weird spiral "cinnamon roll" like objects when viewed in 2D. Their general locations are always correct, but for some reason not all pixels of the cell are assigned to the mask.

This glitch appears to happen downstream of dynamics.compute_masks(). Digging a bit deeper, I've found this effect is sensitive to slightly biasing the values of dP prior to running dynamics.follow_flows(). Indeed, if I slightly reduce the values of dP to 90%, the glitch appears to resolve. If I push dP all the way down to 10%, I still largely get expected output in the sample (but probably lose a few masks). Conversely, if I bump up dP to 110%, additional glitches appear!

I tried to follow the logic of dynamics.follow_flows() and dynamics.get_masks_torch(), but quickly got lost in the latter, so seeking help from the experts! My thought at this moment is if this might be reflecting some kind of "clipping" effect on the flows?

The following code (and sample data) reproduces the plot below. The pre-computed dP and cellprob are from a much larger volume, which I then cropped to a 100x100x100 pixel cube for this demo. The plot shows a few sub-slices from this sample volume. First row is cellprob, second row is dP[1,...], and the subsequent rows are the mask results of biasing dP before follow_flows().

Sample data via image.sc post:
test_probs.tif
test_flows.tif

Plot code

import numpy as np
import imageio.v3 as iio
import matplotlib.pyplot as plt
import colorcet as cc
import torch
from cellpose.dynamics import follow_flows  # v4.1.1
from cellpose.dynamics import get_masks_torch  # v4.1.1

# Load 3D flows/dP [3 x 100 x 100 x 100] and cellprob [100 x 100 x 100] samples.
# These are crops from a much larger volume processed by a custom CP-SAM model.
dP = iio.imread(r'test_flows.tif')
cellprob = iio.imread(r'test_probs.tif')

# Plot customization.
cmap = cc.cm.glasbey_dark
cmap.set_bad('k')
fig, axes = plt.subplots(6, 5, figsize=(15,18))
z_range = range(38, 51, 3)
y_range = slice(30, 60)
x_range = slice(40, 70)

# Show cellprob & one dP axis for zoomed portion of sample volume.
for j, z in enumerate(z_range):
    axes[0, j].imshow(cellprob[z, y_range, x_range], cmap='gray')
    axes[1, j].imshow(dP[1, z, y_range, x_range], cmap='gray')
axes[0, 0].set_ylabel('cellprob')
axes[1, 0].set_ylabel('dP[1, ...]')

# Factors to tweak flows/dP before dynamics.
fudge_factors = [1.1, 1.0, 0.9, 0.1]

# Run through abbreviated version of dynamics.compute_masks() for each fudge factor and plot.
inds = np.nonzero(cellprob > 0)
device = torch.device('cuda')  # torch.device('cpu')  <- doesn't affect outcome
for i, ff in enumerate(fudge_factors):
    p_final = follow_flows(dP * ff * (cellprob > 0) / 5., inds, 200, device)  # <- note insertion of ff here
    masks = get_masks_torch(p_final.int(), inds, dP.shape[1:], 0.4).astype(np.float32)
    masks[masks == 0] = np.nan
    for j, z in enumerate(z_range):
        axes[i+2, j].imshow(masks[z, y_range, x_range], cmap=cmap, vmin=0, vmax=128)
    axes[i+2, 0].set_ylabel(f'{ff:0.1f} * dP')

# Plot customization.
for ax in axes.flatten():
    ax.set_xticks([])
    ax.set_yticks([])

[mrariden]

We've seen something similar before, which I hoped was fixed. We'll look into this further.

Related: #1398, #1408

[mmuunniiaakk]

Interesting, I missed #1398, it does look similar.

My run of model.eval() that generated the above sample did have a slight anisotropy=(2.0/1.866) with flow3d_smooth=0 ... but I also re-ran the volume with anisotropy=None and reproduced the glitchy mask problem, albeit generally at different mask/cell locations throughout the volume.

For reference, my run of the full volume contained ~1,950 masks, and 10 of them appeared to exhibit this glitch (quickly isolated by using skimage.measure.regionprops() and selecting masks with 'euler_number' < -1). But when I re-ran w/o anisotropy, I got ~50 glitchy masks... so anisotropy might not be the critical factor here?

Thanks for looking into it!

[derekthirstrup]

@mmuunniiaakk @mrariden The anisotropic smoothing PR fixed this issue that I reported in issue 1398. Set Z smoothing=1-2 and XY smoothing=0 resolved this issue. Also set anisotropy=1 to prevent z upsampling with synthetic z slices which significantly slows down runtime and does not fix the issue.