Unable to evaluate 3D unet in CLI

[buswinka]

Hello! I am trying to evaluate the 3D UNet on a plant segmentation datasets via the CLI with this command: https://github.com/security
python -m omnipose --omni --dir ./ --dim 3 --z_axis 0 --all_channels --use_gpu --verbose

The input images are greyscale 3D tifs of shape [Z, X, Y]. The CLI seems to evaluate a 2D model on each slice, instead of 3D model, then fails here:

Traceback (most recent call last):
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/site-packages/omnipose/__main__.py", line 3, in <module>
    main(omni_CLI=True)
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/__main__.py", line 344, in main
    out = model.eval(image, channels=channels, diameter=diameter, rescale = args.rescale,
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/models.py", line 300, in eval
    masks, flows, styles = self.cp.eval(x,
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/models.py", line 706, in eval
    masks, styles, dP, cellprob, p, bd, tr = self._run_cp(x,
  File "/home/chris/anaconda3/envs/omnipose/lib/python3.10/site-packages/cellpose/models.py", line 900, in _run_cp
    p = np.array(p)
ValueError: could not broadcast input array from shape (2,512,512) into shape (2,)

Is there a better way to evaluate the 3D model from the CLI? Thank you!

[kevinjohncutler]

@buswinka That is an oversight on my part, I may not have validated the use of 3D models via CLI. That command looks right to me with --dim 3, so something is just going wrong with passing parameters in main.py. I will work on a fix tonight.

[buswinka]

Ayy not a problem. I got it running in 3D via the api. Not worth losing sleep over. Its good to double check! Thank you!