Add SemanticSamTrainer

[anwai98]

@constantinpape Here is the trainer for semantic segmentation using SAM. Let me know if this aligns with what we discussed.

[constantinpape]

Looks good to me, I couldn't spot any issues.

[anwai98]

Thanks @constantinpape. I tested this in a 2d dataset, looks like it's doing the job as expected. This is GTG from my side now (only pending a few minor discussion in the evaluation PR)

[constantinpape]

We should revisit this with multi-class segmentation in mind, and potentially prefer the low-res masks.

[anwai98]

Hi @constantinpape,

Looks like the mutli-class semantic segmentation works as expected now (atleast from the first looks of the Tensorboard logs). I am not a big fan of the workarounds I had to apply to make this work, but maybe we find a better way to make things work in a much more modular setup. Let me know if you spot something. We can discuss together tomorrow the details anyways.

ADDITION: I added the support for an added loss function (cross entropy) over the logits (between the low_res_masks returned by the model and the downscaled version of the ground-truth). Looks like it's working as expected, and converges a bit faster compared to just dice over masks.

[constantinpape]

Looks good, but let's remove the GT downsampling. (I think we don't need it anyore)

[anwai98]

@constantinpape,

I've removed the downscaling of masks. Should be GTG now. Thanks!

PS. Tested it on a quick training as well