At the moment there is a hard condition for the segmentation to be 3D. Should be investigated if or how it works for images with more or less dimensions.