[BUG] Crashing during headless processing a folder of 3D z stacks

[derekthirstrup]

Describe the bug
This crash happens consistently segmenting specific 3D images in a folder with cellpose 3.1.0 . The 3d segmentation of the exact same image works fine with the GUI using same settings so this appears to only occur with headless processing mode.

To Reproduce
Steps to reproduce the behavior:
python -m cellpose --verbose --use_gpu --dir \networkdrive\C2_EGFP_zlibcompressed --pretrained_model cyto3 --restore_type denoise_cyto3 --z_axis 0 --stitch_threshold 0.5 --diameter 80 --flow_threshold 0.0 --cellprob_threshold 0.0 --chan 0 --save_tif --no_npy --min_size 10000 --savedir \networkdrive\CPMask_80_CP310_-1_DefaultIntNorm

Run log
cellpose version: 3.1.0
platform: win32
python version: 3.12.7
torch version: 2.5.1

2024-10-31 14:21:45,152 [INFO] 0%| | 0/101 [00:00<?, ?it/s]
2024-10-31 14:21:45,199 [INFO] No cell pixels found.
C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\dynamics.py:508: UserWarning: The use of x.T on tensors of dimension other than 2 to reverse their shape is deprecated and it will throw an error in a future release. Consider x.mT to transpose batches of matrices or x.permute(*torch.arange(x.ndim - 1, -1, -1)) to reverse the dimensions of a tensor. (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\TensorShape.cpp:3687.)
return pt[..., [1, 0]].squeeze().T
2024-10-31 14:21:47,714 [INFO] 41%|#### | 41/101 [00:02<00:03, 16.00it/s]
2024-10-31 14:21:47,714 [INFO] 10%|# | 23/223 [09:54<1:26:13, 25.87s/it]
Traceback (most recent call last):
File "", line 198, in run_module_as_main
File "", line 88, in run_code
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose_main.py", line 358, in
main()
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose_main.py", line 205, in main
out = model.eval(
^^^^^^^^^^^
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\denoise.py", line 583, in eval
masks, flows, styles = self.cp.eval(
^^^^^^^^^^^^^
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\models.py", line 534, in eval
masks = self._compute_masks(x.shape, dP, cellprob, flow_threshold=flow_threshold,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\models.py", line 620, in _compute_masks
outputs = dynamics.resize_and_compute_masks(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\dynamics.py", line 839, in resize_and_compute_masks
mask = compute_masks(dP, cellprob, niter=niter,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\dynamics.py", line 906, in compute_masks
mask = get_masks_torch(p_final, inds, dP.shape[1:],
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\derekt.conda\envs\cellpose3.1.0\Lib\site-packages\cellpose\dynamics.py", line 757, in get_masks_torch
coo = torch.sparse_coo_tensor(pt, torch.ones(pt.shape[1], device=pt.device, dtype=torch.int),
~~~~~~~~^^^
IndexError: tuple index out of range

[derekthirstrup]

Bug Summary: IndexError and Deprecation Warning in dynamics.py
Overview
When using the latest version of Cellpose 3.1.0, I encountered two primary issues in the dynamics.py file:

Deprecation Warning:

Message:

UserWarning: The use of x.T on tensors of dimension other than 2 to reverse their shape is deprecated and it will throw an error in a future release. Consider x.mT to transpose batches of matrices or x.permute(*torch.arange(x.ndim - 1, -1, -1)) to reverse the dimensions of a tensor.
Cause: Usage of .T (transpose) on tensors with dimensions other than 2, which is deprecated in PyTorch and slated to raise errors in future releases.

IndexError: tuple index out of range
Cause: Attempting to access pt.shape[1] in the get_masks_torch function when pt does not have a second dimension, leading to an out-of-range index error.
Affected File
cellpose/dynamics.py
Affected Sections and Specific Changes
Function: steps_interp

Issue Addressed: Deprecation Warning due to the use of .T on tensors with more than 2 dimensions.
Original Code:

if ndim==3:
return pt[..., [2, 1, 0]].squeeze().T
else:
return pt[..., [1, 0]].squeeze().T

Modified Code:

if ndim == 3:
pt = pt[..., [2, 1, 0]].squeeze()
if pt.dim() == 1:
pt = pt.unsqueeze(1)
else:
pt = pt.permute(1, 0)
return pt
else:
pt = pt[..., [1, 0]].squeeze()
if pt.dim() == 1:
pt = pt.unsqueeze(1)
else:
pt = pt.mT
return pt

Reason for Change:
For 3D Tensors (ndim == 3): Replaced .T with .permute(1, 0) to correctly handle tensor transposition without triggering deprecation warnings.
For 2D Tensors: Replaced .T with .mT (matrix transpose) as .mT is the recommended method for transposing 2D tensors.
Added Dimension Checks: Ensured that tensors are at least 2D before applying transpose operations to prevent future errors.
Function: get_masks_torch

Issue Addressed: IndexError due to accessing pt.shape[1] when pt may not have a second dimension.

Original Code:
coo = torch.sparse_coo_tensor(pt, torch.ones(pt.shape[1], device=pt.device, dtype=torch.int),
shape)

Modified Code:

Ensure 'pt' is 2D before creating the sparse tensor

if pt.dim() == 1:
pt = pt.unsqueeze(1)

Create sparse COO tensor with error handling

try:
coo = torch.sparse_coo_tensor(pt, torch.ones(pt.shape[1], device=pt.device, dtype=torch.int),
shape)
except IndexError as e:
dynamics_logger.error(f"Error creating sparse tensor: {e}")
raise

Reason for Change:
Dimension Check: Before accessing pt.shape[1], the code now checks if pt is 1D and reshapes it to 2D using unsqueeze(1) if necessary. This prevents the IndexError by ensuring that pt always has at least two dimensions.
Error Handling: Wrapped the creation of the sparse tensor in a try-except block to catch and log any IndexError that might still occur, providing clearer diagnostics.

Function: get_masks_torch (Additional Safety)

Issue Addressed: Potential failure when no seeds are found, which could lead to unexpected behavior or errors in subsequent processing.
Original Code:

seeds1 = torch.nonzero((h1 - hmax1 > -1e-6) * (h1 > 10))
npts = h1[tuple(seeds1.T)]
isort1 = npts.argsort()
seeds1 = seeds1[isort1]
Modified Code:
python
Copy code
seeds1 = torch.nonzero((h1 - hmax1 > -1e-6) * (h1 > 10))
del hmax1
if seeds1.numel() == 0:
dynamics_logger.warning("No seeds found in get_masks_torch.")
return np.zeros(shape0, dtype="uint16")
npts = h1[tuple(seeds1.T)]
isort1 = npts.argsort()
seeds1 = seeds1[isort1]

Reason for Change:
Seed Presence Check: Added a condition to check if any seeds are found (seeds1.numel() == 0). If no seeds are detected, the function logs a warning and returns an empty mask. This prevents downstream errors that would occur from processing an empty seed list.
Function: get_masks_torch (Sparse Tensor Creation)

Issue Addressed: Ensuring that the creation of the sparse tensor does not fail due to unexpected tensor dimensions.
Modified Code:

Ensure 'pt' is 2D before creating the sparse tensor

if pt.dim() == 1:
pt = pt.unsqueeze(1)

Create sparse COO tensor with error handling

try:
coo = torch.sparse_coo_tensor(pt, torch.ones(pt.shape[1], device=pt.device, dtype=torch.int),
shape)
except IndexError as e:
dynamics_logger.error(f"Error creating sparse tensor: {e}")
raise

Reason for Change:
Dimension Assurance: Ensured that pt is at least 2D to safely access pt.shape[1].
Error Logging: Implemented error logging to capture and relay any unexpected issues during sparse tensor creation.
Summary of Reasons for All Changes
Addressing Deprecated Operations:

Replaced deprecated .T operations with .mT and .permute() to comply with the latest PyTorch standards and prevent future errors.

Preventing Index Errors:
Added dimension checks and reshaping (unsqueeze) to ensure tensors have the expected number of dimensions before accessing specific indices.
Implemented error handling (try-except) to catch and log potential IndexError occurrences during tensor operations.

Enhancing Robustness:
Introduced checks to handle cases where no seeds are found, thereby avoiding processing on empty or invalid data structures. Ensured that the code gracefully handles edge cases, such as empty masks or large images that might exceed memory constraints. I have a version of the dynamics.py file with the bug fixes for these issues if you would like me to make a PR.

[carsen-stringer]

thanks, I'd really appreciate a PR for this, sorry for the issues

[derekthirstrup]

@carsen-stringer I made a PR and included a minimal change dynamics.py version that addresses the index error as well as an optimized dynamics.py with additional performance optimizations.

[jazberna1]

Hello,

I encountered the same error as described above on a particular 3D z stack, that works on the GUI but not on the command line.

Thanks for the PR!

Jorge