Output class probabilites for segmetation tasks

[geoffreyjdawson]

Currently when we performing inference on a segmentation task we can only output the most probable class. It would be good to have an option to also output the probabilities for each class as separate bands i.e. not perform the argmax.

terratorch/terratorch/tasks/segmentation_tasks.py

Line 344 in 30dfdf1

y_hat = y_hat.argmax(dim=1)

To do this would also mean that we would need an option for a multi-band output here

https://github.com/IBM/terratorch/blob/30dfdf15716edcc9295559aec4327dee62b0d956/terratorch/cli_tools.py#L87C1-L103C1

I have tried this and it works on a single tile, but not with tiled inference.

[Joao-L-S-Almeida]

Is PR #393 solving your issue, @geoffreyjdawson ?
Is it something interesting to include this feature also for classification ?

[romeokienzler]

@geoffreyjdawson agreed to review by EOD

[geoffreyjdawson]

I tested it and for the dataset I am using, the number of output channels is the size of the image (e.g. 512x512x512) instead of the number of classes (e.g. 2x215x512)

[romeokienzler]

merging #393

@geoffreyjdawson please test on version 0.99.9 once released