[BUG] 3D GUI: selecting gradZ view breaks Z navigation (IndexError in gui.py:update_plot, followed by pyqtgraph data.shape[2] error)

[derekthirstrup]

Environment

Cellpose: 4.0.6

OS: Windows (win32)

Python: 3.13.0

Torch: 2.7.1+cu128 (CUDA detected and used)

Repro image: 3D TIFF, 17 Z planes, shape reported by GUI: (17, 58, 45, 3)

What I did (repro steps):

1. Launched cellpose --Zstack and load a Z-stack (17 planes).

2. Ran fine tuned segmentation model (GPU).

3. In the viewer, switch the display to gradZ.

4. Move the Z slider.

Expected behavior:

The 2D plane at the current Z to render correctly for gradZ, and the Z slider to work normally.

What actually happened

As soon as gradZ is selected and I move the Z slider, the GUI throws repeated IndexErrors from gui.py:update_plot and then pyqtgraph rendering errors:

File ".../cellpose/gui/gui.py", line 1388, in update_plot image = self.flows[self.view - 1][self.currentZ] ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^ IndexError: index 6 is out of bounds for axis 0 with size 2

Then, repeatedly while moving Z:

IndexError: index {2..16} is out of bounds for axis 0 with size 2/3

TypeError: data.shape[2] must be <= 4 (Representative snippets from the session:)

IndexError: index 6 is out of bounds for axis 0 with size 2

IndexError: index 8 is out of bounds for axis 0 with size 3

TypeError: data.shape[2] must be <= 4

Minimal traceback excerpt

File ".../cellpose/gui/gui3d.py", line 568, in update_plot super().update_plot() File ".../cellpose/gui/gui.py", line 1388, in update_plot image = self.flows[self.view - 1][self.currentZ] IndexError: index 6 is out of bounds for axis 0 with size 2

Likely root cause (dimension/axis mix-up) In 3D, the flows array for gradient views appears to be stored with a components-first layout, e.g. (C, Z, Y, X) where C ∈ {2,3} (components such as dZ/dY/dX or dY/dX).

update_plot assumes Z is axis 0 and does:

image = self.flows[self.view - 1][self.currentZ] which indexes axis 0 with currentZ. If axis 0 is actually the component axis (size 2 or 3), this produces the observed:

IndexError: ... axis 0 with size 2/3

After the exception, the code path sometimes hands a 3D array with shape like (Y, X, Z) to pyqtgraph.ImageItem, which then complains:

TypeError: data.shape[2] must be <= 4 because pyqtgraph interprets the last dimension as color channels and only supports up to 4 (RGBA). Passing Z as the last dimension violates that constraint.

In short: gradZ view is using the wrong axis for Z when slicing flows, and the fallback display path can end up feeding a (Y, X, Z) array to pyqtgraph as if it were an RGB(A) image.

Proposed fix (robust Z-slicing + channel guard) Two parts:

Index Z along the correct axis for flows, not hard-coded axis 0. A small, backward-compatible guard in gui.py:update_plot can detect the Z axis by matching the number of planes in the stack, and handle component-first vs. Z-first layouts.

Guarantee a 2D plane (or <=4 channels) is sent to ImageItem. If after slicing Z the image is still 3D with more than 4 channels (e.g., (Y, X, Z) or (C, Y, X)), reduce to a single component (for gradZ) or compute a magnitude; avoid handing (Y, X, Z) to pyqtgraph.

[derekthirstrup]

I have fixed this 3D indexing bug in gui.py and here is a link to the PR. https://github.com/MouseLand/cellpose/pull/1295/commits/b57d960f6a3c00e5af6f988d1380399d8639b897