Feature request: allow segmentation to work when axes are reordered, e.g. orthogonal view

[psobolewskiPhD]

So right now it always runs on the 0 axis, but it would be cool if it was aware of viewer.dims.displayed and used the non-displayed axis, this way you could change the axis order in napari and segment along a different orthogonal view.

[constantinpape]

I think that makes sense, but it's a bit tricky. We pre-compute the embeddings for 3D segmentation, otherwise the segmentation is not interactive. So when switching the axis order the emebddings also need to be recomputed (which can take quite long on a CPU). I won't have time to do this myself in the near future, but I am happy to provide you with more concrete pointers in case you want to give it a go.

[psobolewskiPhD]

Oh, I didn't realize the embeddings would need to be recomputed! I sort of assumed they'd be invariant I guess.

[constantinpape]

No, the model processes each image slice by slice, with the embedding being computed for each slice. So if you want to do any reslicing of the data, including axis permutations, they need to be re-computed.

That being said, for the axis orders we would just need three different versions of the embeddings, so at least on a GPU pre-computing these would not be such a big overhead.

[psobolewskiPhD]

Thanks for the explanation. Long term would be nice as an option -- the embeddings take up a non-trivial amount of disk space, something the user I'm working with is concerned about. Short term, maybe this a documentation point? I assume if I used np.moveaxis to reorder my image array before computing the embeddings I could use that? With the caveat that those embeddings would be just for that order of course.

[constantinpape]

Long term would be nice as an option

:+1: I will leave this issue open.

the embeddings take up a non-trivial amount of disk space, something the user I'm working with is concerned about.

It's probably best to remove them once a volume is done. And if the user has a GPU it's probably not necessary to cache them at all. Caching the embeddings to disc is optional, and mostly for users without a GPU to avoid long re-computation times.

I assume if I used np.moveaxis to reorder my image array before computing the embeddings I could use that? With the caveat that those embeddings would be just for that order of course.

Yes, exactly. That would work.