Change 2-class model to 3-class model

[hermancollin]

At the moment, the model has 2 output channels for the axon and myelin masks. This formulation of the problem is not ideal because the classes are not mutually exclusive and the model gives lots of false positives for the background class. This could be mitigated by having a separate output channel for the background and using one-hot encoding for the loss/metrics computations.

The explanation related to this for the ivadomed implementation can be found in this comment

For reference, the current predictions look like this

[hermancollin]

After some preliminary tests on the loader (see updated notebook), I found it was tricky to merge the labels. Fortunately, we have the axonmyelin-manual labels included in our datasets. Note that we need to add these 2 transforms to the labels:

NormalizeIntensityd(keys="label", subtrahend=0, divisor=127, nonzero=True),
AsDiscreted(keys="label", rounding='torchrounding')

Otherwise, one-hot encoding will fail. The first transform will change the [0, 127, 255] values to [0, 1.something, 2.something], which we can then round with the latter transform to get discrete values.