Adaptation to Multiclass Segmentation

[vminin2000]

Hi! I hope to leave a public record here to help others that are struggling to convert binary semantic segmentation code to multi-class semantic segmentation code. I reached out to you and have done other research as to how to go about adapting your code and I have come to the conclusion that I must change the deeplabv3_plus function definition's last four lines of code to be this

x = Conv2D(4,(1, 1))(x) ## 4 is the number of classes
x = Activation("softmax")(x) ## The activation function has been changed from sigmoid to softmax
model = Model(inputs, x)
return model

Then the loss function must be changed as well, I found recommendations to change it to CrossEntropyLoss. The adaptation occurs in the train.py file on line 99

criterion = torch.nn.CrossEntropyLoss()
model.compile(loss=dice_loss, optimizer=Adam(lr), metrics=[dice_coef, iou, Recall(), Precision()])

I believe the metrics must be changed also to Jaccard Score (IOU), and Dice Score.

I am receiving errors such as shapes being incorrect once I attempt to train the model and I am not sure what tensor object to insert in the loss function. If there is any advice you can give that'd be much appreciated. I know by myself and many others. Thanks!