From ef8d4436abcf7d2de2c1e8ac15e1d43306fad873 Mon Sep 17 00:00:00 2001 From: andy-s-ding Date: Mon, 2 Dec 2024 13:29:09 -0500 Subject: [PATCH] Added template framework for applying transforms to distance maps nnUNet currently does not have a framework for loading in distance maps that were calculated offline, though I have a version of nnUNet with a modified dataloader that can do this. For full functionality, distance maps should also be cropped and transformed in a similar fashion to the image and segmentation data. The easiest way to do this is to add another key in the data dictionary called 'dist_map'. If the data has this key, then a transform can call _apply_to_dist_map. Modifying the BasicTransform class is the simplest way to add this functionality to other transforms. --- batchgeneratorsv2/transforms/base/basic_transform.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/batchgeneratorsv2/transforms/base/basic_transform.py b/batchgeneratorsv2/transforms/base/basic_transform.py index 1db856b..7cd1064 100644 --- a/batchgeneratorsv2/transforms/base/basic_transform.py +++ b/batchgeneratorsv2/transforms/base/basic_transform.py @@ -26,6 +26,9 @@ def apply(self, data_dict, **params): if data_dict.get('segmentation') is not None: data_dict['segmentation'] = self._apply_to_segmentation(data_dict['segmentation'], **params) + + if data_dict.get('dist_map') is not None: + data_dict['dist_map'] = self._apply_to_dist_map(data_dict['dist_map'], **params) if data_dict.get('keypoints') is not None: data_dict['keypoints'] = self._apply_to_keypoints(data_dict['keypoints'], **params) @@ -44,6 +47,9 @@ def _apply_to_regr_target(self, regression_target, **params) -> torch.Tensor: def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor: pass + def _apply_to_dist_map(self, dist_map: torch.Tensor, **params) -> torch.Tensor: + pass + def _apply_to_keypoints(self, keypoints, **params): pass @@ -74,4 +80,4 @@ def apply(self, data_dict: dict, **params) -> dict: if __name__ == '__main__': - pass \ No newline at end of file + pass