[BUG] Flow generation is perturbed by neighboring cells #1486

Before you fill out this form:

Did you review the FAQ?

• Yes, I reviewed the FAQ on the cellpose ReadTheDocs

Did you look through previous (open AND closed) issues posted on GH?

• Yes, I searched for related previous GH issues.

Now fill this form out completely:

Describe the bug

Hello, first I wanted to thank you for your great work with cellpose.

I was trying to train a new model using hand drawn 3D cells, and upon flows generation I realized that the flows of a cell were perturbed by neighbouring cells, while the flows of a single cell were fine. This leads to the interface behind the 2 cells being shifted inside one cell when the flows are followed.
Upon checking the code I realized that the gradients were computed with the function _extend_centers_gpu (in dynamics), using the difference of heat of the neighbors without taking into account if the 2 pixels belonged to the same cell. It is a bug present in 2D and 3D (though it seems to have a much more dramatic effect in 3D with tilted cells).

Please find attached a single 3D slice of the original masks and the flows generated by the current version of cellpose as well as the masks regenerated following the flows.

To Reproduce

Steps to reproduce the behavior:

from cellpose.dynamics import masks_to_flows_gpu_3d, follow_flows, get_masks_torch
import tifffile
import numpy as np
masks = tifffile.imread("masks.tif")
inds = np.nonzero(masks>0)
flows = masks_to_flows_gpu_3d(masks)
p = follow_flows(flows, inds)
regenerated_masks = get_masks_torch(p.int(), inds, masks.shape)
tifffile.imwrite("masks_regenerated.tif", regenerated_masks)
tifffile.imwrite("flows.tif", flows)

Screenshots

input masks | ZYX flows | reconstructed masks from flows

please find here the actual 3D input data I used : masks.tif