adding focal loss, distance map dice loss, loading distance maps

[andy-s-ding]

Added a FocalLoss class that also works in compound loss functions. nnUNet also currently does not have a framework for loading in distance maps that were calculated offline. I propose a framework of calculating distance maps from segmentations in the preprocessed folder designated by the plans json. These maps are saved as <case_num>_dist.npy. In doing so, the plans json does not need to necessarily be modified. I have also made a Dataset class and a DataLoader class that is able to load these distance maps during training. Furthermore, I have added custom transform classes that transform distance maps commensurately with image data and segmentations. To make this work, I submitted a pull request in batchgeneratorsv2 to modify the BasicTransform class slightly.

[andy-s-ding]

I also added compute_distance_maps.py in the preprocessing folder. One major TODO is to integrate distance map computation into the preprocessing pipeline. For now, I am taking the _seg.npy files after preprocessing, computing the distance maps, and saving as _dist.npy files.

Currently training distance-weighted loss functions right now and can confirm that epoch time is essentially the same (56.7 seconds per epoch for CE + DistDiceLoss vs. 51.2 seconds per epoch for CE + DiceLoss) on an RTX 3090 GPU.

Again, based on my edits, this is contingent on BasicTransform identifying 'dist_map as a key in the data_dict. I have submitted a pull request for this change in batchgeneratorsv2 (MIC-DKFZ/batchgeneratorsv2#10)

[FabianIsensee]

Hey,
thanks a lot for the effort and for putting this together. The FocalLoss addition is definitely useful — would be great to have that as a standalone PR so we can merge it cleanly.

As for the distance map pipeline: I totally get the motivation, and it’s cool to see you're already training with it. That said, it introduces quite a bit of additional complexity across preprocessing, transforms, and data loading — and without a clear or widely requested use case, I’m hesitant to bring that into the core codebase. The current pipelines are already quite complex, and maintaining forward compatibility (plus answering future issues) is not something I want to commit to unless the benefit is substantial.

So my suggestion would be: let’s keep the distance map logic as a custom extension for now, outside the main repo. If there’s more traction or a clear need in the future, we can definitely revisit it.
Best,
Fabian