Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 5 commits into
base: master
Choose a base branch
from

Conversation

andy-s-ding
Copy link

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 andy-s-ding marked this pull request as draft December 2, 2024 18:37
@FabianIsensee FabianIsensee self-assigned this Dec 3, 2024
@andy-s-ding
Copy link
Author

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)

@andy-s-ding andy-s-ding marked this pull request as ready for review December 3, 2024 17:31
@andy-s-ding andy-s-ding marked this pull request as draft December 3, 2024 17:43
@andy-s-ding andy-s-ding marked this pull request as ready for review December 3, 2024 18:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants