[BUG] Shape mismatch between masks and output when segmenting z-stack slice-by-slice

[ElpadoCan]

Describe the bug

When segmenting a 3D z-stack volume with shape (24, 2048, 2048) and diameter = 20 the image is correctly resized to (24, 3072, 3072) before running _run_net. However, if stitch_threshold > 0.0 and do_3D = False, the outputs variable inside _compute_masks are not resized to (24, 3072, 3072) while masks is initialized with shape (24, 3072, 3072). Therefore when cellpose tries masks[i] = outputs there is a shape mismatch.

More specifically, I get this error

Traceback (most recent call last):
  File "/home/schmollerlab/Cell_ACDC/cellacdc/test_segm_model.py", line 135, in <module>
    lab = core.segm_model_segment(
          ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/schmollerlab/Cell_ACDC/cellacdc/core.py", line 1934, in segm_model_segment
    lab = model.segment(image, **model_kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/schmollerlab/Cell_ACDC/cellacdc/models/cellpose_v4/acdcSegment.py", line 327, in segment
    labs = self.eval_loop(
           ^^^^^^^^^^^^^^^
  File "/home/schmollerlab/Cell_ACDC/cellacdc/models/_cellpose_base/acdcSegment.py", line 373, in eval_loop
    labels = self._eval(images, **eval_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/schmollerlab/Cell_ACDC/cellacdc/models/_cellpose_base/acdcSegment.py", line 151, in _eval
    out, removed_kwargs = myutils.try_kwargs(
                          ^^^^^^^^^^^^^^^^^^^
  File "/home/schmollerlab/Cell_ACDC/cellacdc/myutils.py", line 5168, in try_kwargs
    return func(*args, **kwargs), removed_kwargs
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/schmollerlab/miniforge3/envs/acdc/lib/python3.12/site-packages/cellpose/models.py", line 341, in eval
    max_size_fraction=max_size_fraction, niter=niter,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/schmollerlab/miniforge3/envs/acdc/lib/python3.12/site-packages/cellpose/models.py", line 542, in _compute_masks
    if stitch_threshold > 0 and nimg > 1:
            ^^^^^^^^
ValueError: could not broadcast input array from shape (2048,2048) into shape (3072,3072)

To Reproduce
Steps to reproduce the behavior:

Run the following script

import skimage.io

from cellpose.models import CellposeModel

zstack_image_filepath = r'path_to_zstack_data'

zstack_data = skimage.io.imread(zstack_image_filepath)

print(zstack_data.shape)

cp_model = CellposeModel(
    pretrained_model='cpsam', 
    gpu=True,
)

eval_kwargs = {
    'diameter': 20.0, 
    'flow_threshold': None, 
    'stitch_threshold': 0.5, 
    'do_3D': False, 
    'anisotropy': 1.0, 
    'max_size_fraction': 0.4, 
    'invert': False, 
    'flow3D_smooth': 0.0, 
    'niter': None, 
    'augment': False, 
    'tile_overlap': 0.1, 
    'bsize': 256, 
    'min_size': 15, 
    'cellprob_threshold': 0.0, 
    'normalize': {
        'lowhigh': None, 
        'percentile': (1.0, 99.0), 
        'norm3D': False, 
        'tile_norm_blocksize': 0
    }, 
    'batch_size': 8, 
    'channel_axis': None, 
    'z_axis': 0
}

masks, flows, styles = cp_model.eval(zstack_data, **eval_kwargs)

I did some digging and I found out that the keyword argument resize in cellpose.dynamics.resize_and_compute_masks is deprecated (although the user does not have the possibility to change that kwarg) and not used, even though in my case, resize = [3072, 3072] which sounds like the correct value to resize outputs before inserting into masks.

[dfgdgdfgd]

Hi,
unfortunately, I don't have any programming knowledge, but I asked ChatGPT for help and was able to successfully integrate the following path into my CellACDC environment:

Minimal patch (honour resize in 2D):
Search for the function resize_and_compute_masks in dynamics.py.
Insert this block before the return mask (or replace existing ‘deprecated’ lines with it):

--- begin patch: respect resize for 2D ---

if (not do_3D) and (resize is not None):
try:
from cellpose import transforms
import cv2
Ly, Lx = int(resize[0]), int(resize[1])
if mask.shape[-2:] != (Ly, Lx):
mask = transforms.resize_image(
mask, Ly=Ly, Lx=Lx, no_channels=True,
interpolation=cv2.INTER_NEAREST
)
except Exception as e:
dynamics_logger.warning(f"resize requested but failed ({e}); continuing without resizing")

--- end patch ---

Now Cellpose4 is working again with CellACDC.....perhaps one of the developers could integrate this patch into the upcoming Cellpose version.

Many thanks!!