How to configure the models and ground truth for multi-class segmentation?

[gokriznastic]

Hey, I am adapting this code to work on a different dataset that was used in LiTS competition last year. Therefore, I have two distinct classes + the background.

Currently, I am making the model predict output with 2 channels, one for each class. I am splitting the ground truth into 2 channels too for comparing with the prediction during loss calculation. At the end of the loop am merging both channels again to make it look like original groundtruth (for visualization purposes).

The problem in this case is that the model seems to be learning one of the classes well but doesn't learn the other.

Can you please provide some pointers as to why this may be happening or is there any better way to configure SegAN for multi-class segmentation.

Thanks.