process_block, mismatched shape in zarr

[Elena983]

When running distributed_eval() from cellpose.contrib.distributed_segmentation on
a 3D Zarr input (channels, H, W), I encountered a shape mismatch error during block processing.

Specifically, output_zarr[tuple(crop)] = segmentation fails due to a shape mismatch between the expected crop shape and the segmentation output.

Env:

Cellpose Version: 4.0.7.dev7+gdf6b944
Python Version: 3.10.13
Platform: Linux
Torch Version: 2.7.1+cpu
Execution Mode: CPU (note: GPU not available)

zarr_overview("input_image.zarr"):

Type: Array
Path: /home/input_image.zarr
Shape: 3 x 640 x 480
Chunk Shape: 1 x 128 x 128
No. of Chunks: 60 (3 x 5 x 4)
Data Type: float64
Compressor: zlib

the code chunk:

result = distributed_segmentation.distributed_eval(
    input_zarr=input_zarr,
    blocksize=(1, 128, 128),
    write_path=str(write_path),
    model_kwargs=dict(gpu=False),
    eval_kwargs=dict(
        diameter=30,
        z_axis=None,
        do_3D=False
    ),
    cluster_kwargs=dict(
        n_workers=1,
        ncpus=32,
        memory_limit="256GB",
        threads_per_worker=1,
        dashboard_address=":0"
    )
)


error:

RUNNING BLOCK:  (0, 0, 0)     REGION:  (slice(np.int64(0), np.int64(3), None), slice(np.int64(0), np.int64(188), None), slice(np.int64(0), np.int64(188), None))
Cluster dashboard link: http://192.168.22.24:35487/status

2025-08-04 16:10:12,437 - distributed.worker - ERROR - Compute Failed
Key:       process_block-31d3c011637b357384f2764f074c1bbd
Exception: ValueError("parameter 'value': expected array with shape (1, 128, 188), got (1, 128)")

Traceback:
  File ".../cellpose/contrib/distributed_segmentation.py", line 500, in process_block
    output_zarr[tuple(crop)] = segmentation
  File ".../zarr/core.py", line 1447, in __setitem__
    self.set_orthogonal_selection(pure_selection, value, fields=fields)
  File ".../zarr/core.py", line 1636, in set_orthogonal_selection
    self._set_selection(indexer, value, fields=fields)
  File ".../zarr/core.py", line 1962, in _set_selection
    check_array_shape("value", value, sel_shape)
  File ".../zarr/util.py", line 561, in check_array_shape
    raise ValueError(...)

Thanks