How to train multi class segmentation?

[huyu398]

I want to train CapsNetR3 with n_class=3 and shape=(160,160), but Length(num_classes=n_class, seg=True, name='out_seg')(seg_caps) returns a tensor with shape of (?, 160, 160, 1).
Maybe, it should return a tensor with shape of (?, 160, 160, 3).
Therefore, Length.compute_output_shape returns (None, 160, 160, 3).
Please tell me a solution to this problem.

appendix

Keras: ver. 2.2.4
tensorflow-gpu: ver. 1.13.1
Input Image: (160, 160, 3) 160x160 pixels image with 3 channel
Output Image: (160, 160, 3) 160x160 pixels image with 3 classes
Executing Code: train_model, eval_model, manipulate_model = CapsNetR3((160,160,3), 3)

Thank you.

[msseibel]

As far as I understand capsule networks, seg_caps should have as many capsules as there are different classses i.e. 3.
Maybe the Length layer was only tested for 1 class. Especially it seems like the decoder was not designed to take more than 1 class, since the original digit capsule is flattened to num_classes*num_atoms.

So you can try to rewrite the Length layer and to remove the assert.
I have tried segcaps only with grayscale images for binary segmentation. If you can achieve acceptable results with RGB images, please tell about it.