[BUG] distributed segmentation produces discontinuous masks at grid boundaries

[stanfish06]

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.

Describe the bug

Hi! I tried distributed segmentation to segment nuclei, and I found that nuclear masks at tile/block boundaries are being split up. I'm using the code in this pr #1334 to run distributed segmentation with slurm. I also attached my script below. I wonder if I'm doing something wrong in my configuration, or if this is a known behavior? Thank you so much for your time.

run_cellpose_distributive.py

[erjel]

Hi,

thank you for testing the PR!

A quick question: Does your code print the done output?

For me it looks like the code exited before the final relabeling step:

cellpose/cellpose/contrib/distributed_segmentation.py

Lines 776 to 781 in 70c0d07

	relabeled = dask.array.map_blocks(
	lambda block: np.load(new_labeling_path)[block],
	segmentation_da,
	dtype=np.uint32,
	chunks=segmentation_da.chunks,
	)

In my tests, scaling the cluster resources down for the relabeling step did not work properly with the existing code. I am not a dask expert, so that's why I took the brute-force route: modified cellpose/contrib/distributed_segmentation.py and commented (for now) all down-scaling lines in distributed_eval. Our cluster was not fully occupied at the time, so I could afford the additional resource requirements.

If you trust the PR code, you can directly fetch the PR in your local cellpose git repository:

git fetch origin pull/1334/head:pr1334
git switch pr1334

this would allow you to directly use my changes to distributed_segmentation.py with your already existing script.

Best wishes,
Eric

PS: It looks like cellpose is working better on your sample than on the example I used, would it be possible to publicly share (a max 1024x1024x1024 subset of) it so that I can use it as part of the cluster test script and documentation?

[stanfish06]

Dear Eric,

Thanks for your quick response! To answer your questions:

A quick question: Does your code print the done output?

Yes it did. Also, I ran my script based on your commit on Oct 1st (see output below). I installed your commit via pip install -e ..

(hpc) [zyyu@gl-login3 cellpose]$ git log -1
commit 4c8e8a9 (HEAD -> dev, slurm)
Merge: d8ca899 0a3de7b
Author: Eric Jelli 36883188+erjel@users.noreply.github.com
Date: Wed Oct 1 11:24:23 2025 +0200
Merge branch 'MouseLand:main' into distributed

I could also see the debug messages you added before, and I had never touched your edits in the distributed_segmentation.py

print('Use dask array to relabel segmentations')
segmentation_da = dask.array.from_zarr(temp_zarr)
print('Map dask to blocks')
relabeled = dask.array.map_blocks(
    lambda block: np.load(new_labeling_path)[block],
    segmentation_da,
    dtype=np.uint32,
    chunks=segmentation_da.chunks,
)
print("Write output")
dask.array.to_zarr(relabeled, write_path, overwrite=True)
# TODO(erjel): Scale cluster down again?
print("Merge bboxes")
merged_boxes = merge_all_boxes(boxes, new_labeling[box_ids])
print("Merge done")
return zarr.open(write_path, mode='r'), merged_boxes

I just noticed that you have an update today, so I'll try that later. Speaking of scaling down cluster resources for relabeling, I didn't notice that. Maybe I could try to uncomment those lines and see if it works on my University's cluster.

PS: It looks like cellpose is working better on your sample than on the example I used, would it be possible to publicly share (a max 1024x1024x1024 subset of) it so that I can use it as part of the cluster test script and documentation?

I can certainly share some data with you for testing as I also want to make this work, but I don't have image stack with >1000 zs because in my current lab, we study 2D human in vitro models, which is kind of flat. I can send you images of medium sizes, like 221x2835x2450, if you want. I attached an example below in case you want to take a look.
The xy view of nuclear DAPI:

The ortho view:

Best,
Stan

[erjel]

I changed #1334 a little so that it uses your example data by default.

Unfortunately, I cannot reproduce the grid boundaries artifacts:

Maybe you can try to run the cluster_script.py on your system, as well so that we can compare the results ...

[stanfish06]

Hi,

I just tried your updated pr and used your cluster_script.py, and I couldn't reproduce the issue as well. It works perfectly fine now. I'll close the issue. Thank you so much for the help! I hope this pr can get merged soon.