search

qubvel / segmentation_models

Segmentation models with pretrained backbones. Keras and TensorFlow Keras.
MIT License
4.74k stars 1.03k forks source link

Multi class segmentation (choosing more than 2 classes) #346

Open NProdanova opened 4 years ago

NProdanova commented 4 years ago

Thanks for the repo! When I add more classes for example like this: CLASSES = ['car', 'pedestrian','building', 'pavement'] I get an error: ValueError: Dimensions must be equal, but are 5 and 3 for 'loss/softmax_loss/mul_4' (op: 'Mul') with input shapes: [5], [3]. How can I add more classes? It seems like it is coming from the loss function, which one would be the right to use? Thanks in advance!

JordanMakesMaps commented 4 years ago

That looks like it might be an issue with shape of the data you're feeding to the model for training? Are the masks one-hot-encoded? If so, is there shape consistent with the number of classes you're training for?

Also, is this the same as the number of classes in the final layer of your model?

j-3-1415 commented 2 years ago

Not sure if this is still relevant, but this error comes from the fact that the classes have assigned hard-coded weights in the dice loss function.

If you change dice_loss = sm.losses.DiceLoss(class_weights=np.array([1, 2, 0.5])) to dice_loss = sm.losses.DiceLoss() (or enter an array of weights for each of your classes), it works.