Surface dice loss #45

LorenzLamm · 2024-01-03T14:48:46Z

Added the option to use Surface-Dice as a loss function during training.

Surface-Dice is based on "clDice - a Novel Topology-Preserving Loss Function for Tubular Structure Segmentation"
(https://openaccess.thecvf.com/content/CVPR2021/papers/Shit_clDice_-_A_Novel_Topology-Preserving_Loss_Function_for_Tubular_Structure_CVPR_2021_paper.pdf)

Also fixed some issues for patch extraction (corrected naming), removed wandb tracking (caused dependency issues), fixed bug mentioned in #44, and added printing of training parameter summary.

LorenzLamm · 2024-01-03T14:51:45Z

src/membrain_seg/annotations/extract_patch_cli.py

@@ -21,6 +21,11 @@ def extract_patches(
        help="Path to the folder where extracted patches should be stored. \
            (subdirectories will be created)",
    ),


Dataset token is now readable as well. This helps to distinguish between different datasets, because we may want to apply different loss functions (particularly Surface-Dice) to some datasets, but not to others.

LorenzLamm · 2024-01-03T15:47:27Z

src/membrain_seg/segmentation/cli/train_cli.py

@@ -84,6 +87,22 @@ def train_advanced(
            but also severely increases training time.\
                Pass "True" or "False".',
    ),


surface_dice_tokens: dataset tokens specifying which datasets to apply surface-dice to.

Needs to be passed as separate arguments:
--surface-dice-tokens ds1 --surface-dice-tokens ds2

I did not find a more elegant way with Typer to pass in a list of strings

ah that's a bummer, but I think it's okay. In the future, it might make sense to add support for a glob string or something or directory so that people don't have to write all of the tokens.

LorenzLamm · 2024-01-03T15:48:24Z

src/membrain_seg/segmentation/dataloading/memseg_dataset.py

@@ -102,6 +101,7 @@ def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
            "label": np.expand_dims(self.labels[idx], 0),
        }
        idx_dict = self.transforms(idx_dict)


dataset token is now returned with every train image

LorenzLamm · 2024-01-03T15:48:47Z

src/membrain_seg/segmentation/dataloading/memseg_dataset.py

@@ -190,3 +192,23 @@ def test(self, test_folder: str, num_files: int = 20) -> None:
                    os.path.join(test_folder, f"test_mask_ds2_{i}_group{num_mask}.png"),
                    test_sample["label"][1][0, :, :, num_mask],
                )
+
+


dataset token is defined as first token before 1st underscore

LorenzLamm · 2024-01-03T15:49:49Z

src/membrain_seg/segmentation/networks/unet.py

            deep_supervision=True,
            deep_supr_num=2,
        )
-        ignore_dice_loss = IgnoreLabelDiceCELoss(ignore_label=2, reduction="mean")
+


The following initializes a weighted average between loss functions.
BCE/DICE loss are used by default, Surface-Dice is optionally added

codecov-commenter · 2024-01-03T15:50:37Z

Codecov Report

Attention: 125 lines in your changes are missing coverage. Please review.

Comparison is base (1f9747a) 5.41% compared to head (9f39fc7) 7.73%.

Files	Patch %	Lines
...membrain_seg/segmentation/training/surface_dice.py	23.37%	59 Missing ⚠️
...eg/segmentation/training/training_param_summary.py	0.00%	32 Missing ⚠️
src/membrain_seg/segmentation/networks/unet.py	5.26%	18 Missing ⚠️
...ain_seg/segmentation/dataloading/memseg_dataset.py	0.00%	7 Missing ⚠️
.../membrain_seg/segmentation/training/optim_utils.py	90.00%	3 Missing ⚠️
src/membrain_seg/annotations/merge_corrections.py	0.00%	2 Missing ⚠️
src/membrain_seg/segmentation/cli/train_cli.py	0.00%	2 Missing ⚠️
src/membrain_seg/segmentation/train.py	0.00%	2 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff            @@
##            main     #45      +/-   ##
========================================
+ Coverage   5.41%   7.73%   +2.31%     
========================================
  Files         38      40       +2     
  Lines       1256    1410     +154     
========================================
+ Hits          68     109      +41     
- Misses      1188    1301     +113

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

LorenzLamm · 2024-01-03T15:51:55Z

src/membrain_seg/segmentation/networks/unet.py

+            loss_inclusion_tokens.append(surf_dice_tokens)
+
+        scaled_weights = [entry / sum(weights) for entry in weights]
+


The combined loss function computes the losses only for selected datasets.
BCE/DICE are computed for all datasets, for S-Dice datasets can be custom-chosen

LorenzLamm · 2024-01-03T15:52:22Z

src/membrain_seg/segmentation/networks/unet.py

        output = self.forward(images)
-        loss = self.loss_function(output, labels)
+        loss = self.loss_function(output, labels, ds_label)

        stats_dict = {"train_loss": loss, "train_number": output[0].shape[0]}
        self.training_step_outputs.append(stats_dict)
        self.running_train_acc += (
            masked_accuracy(output[0], labels[0], ignore_label=2.0, threshold_value=0.0)
            * output[0].shape[0]
        )


also log surface-dice during training

LorenzLamm · 2024-01-03T15:53:11Z

src/membrain_seg/segmentation/train.py

+    surf_dice_weight : float, optional
+        Weight for the Surface-Dice loss.
+    surf_dice_tokens : list, optional
+        List of tokens to use for the Surface-Dice loss.

    Returns
    -------
    None
    """


This prints a summary of training parameters before each training run.

LorenzLamm · 2024-01-03T15:54:18Z

src/membrain_seg/segmentation/train.py

@@ -106,7 +133,7 @@ def on_epoch_start(self, trainer, pl_module):
    # Set up the trainer
    trainer = pl.Trainer(
        precision="16-mixed",


wandb logging requires additional dependency & wandb registration. Should be an option in the future, but removed for now

LorenzLamm · 2024-01-03T15:54:48Z

src/membrain_seg/segmentation/training/metric_utils.py

@@ -34,7 +34,7 @@ def masked_accuracy(
    mask = (
        y_gt == ignore_label
        if ignore_label is not None
-        else torch.ones_like(y_gt).bool()
+        else torch.zeros_like(y_gt).bool()


Fixed issue mentioned in #44

LorenzLamm · 2024-01-03T15:55:56Z

src/membrain_seg/segmentation/training/optim_utils.py

@@ -53,7 +53,7 @@ class IgnoreLabelDiceCELoss(_Loss):
    def __init__(
        self,
        ignore_label: int,


Loss functions should be default return non-reduced losses. I.e., for each element in the batch, a single value should be returned.
This way, we can decide whether to apply the respective loss for each element in the batch (depending on the dataset token)

LorenzLamm · 2024-01-03T15:57:38Z

src/membrain_seg/segmentation/training/optim_utils.py

+        return loss
+
+
+class CombinedLoss(_Loss):


Combined loss function computes the weighted averages of all losses, and considers only specified datasets.
In this way, we can choose exactly which losses to apply for which dataset

LorenzLamm · 2024-01-03T15:58:09Z

src/membrain_seg/segmentation/training/surface_dice.py

@@ -0,0 +1,455 @@
+"""


Implementation of Surface-Dice functionalities

Soft skeletonization of the segmentation is achieved by iteratively eroding and dilating the membrane, and keeping track of the differences.
Erosion and Dilation can be achieved by min- and max-pooling, respectively. This makes the function differentiable and we can perform backpropagation through it.

LorenzLamm · 2024-01-03T16:01:34Z

src/membrain_seg/segmentation/training/surface_dice.py

+        skel = skel + F.relu(delta - skel * delta)
+    return skel
+
+


Defined Gaussian kernel for smoothing binary segmentations before skeletonization.
I didn't find another torch function for this, so I implemented it.
This allows computation of smoothing on GPU without moving stuff between devices.

LorenzLamm · 2024-01-03T16:02:27Z

src/membrain_seg/segmentation/training/surface_dice.py

+    filtered_seg = F.conv3d(seg, g_kernel, padding=padding, groups=seg.shape[1])
+    return filtered_seg
+
+


for binary segmentations, we first perform Gaussian smoothing. Otherwise soft skeletons are discontinuous

LorenzLamm · 2024-01-03T16:02:51Z

src/membrain_seg/segmentation/training/surface_dice.py

+    skel_gt = soft_skel(gt_smooth, iter_=iterations)
+    return skel_gt
+
+


Computation of Surface-Dice similar to centerline Dice.

LorenzLamm · 2024-01-03T16:03:24Z

src/membrain_seg/segmentation/training/training_param_summary.py

@@ -0,0 +1,117 @@
+def print_training_parameters(


I thought it might be useful to print training parameters before each training run. Maybe makes sense to make this optional?

LorenzLamm · 2024-01-03T16:05:48Z

tests/membrain_seg/training/test_optim_utils.py

Adjusted test function to work with new loss definitions.
I should cover more code with the tests I guess :/

Thanks for updating. It would be nice to increase the test coverage over time, but it's fine for this PR

kevinyamauchi

This looks good to me! I have some minor comments below. Based on the conversation on zulip with @alisterburt , it sounds like we are in agreement that this loss function should live in membrain-seg. I think you can merge after you address the minor comments. Thanks, @LorenzLamm !

kevinyamauchi · 2024-01-20T15:40:55Z

src/membrain_seg/segmentation/training/surface_dice.py

+
+    Parameters
+    ----------
+    data : torch.Tensor


It looks like the code assumes a certain shape for the data array. It would be nice to write the expected axis order in the docstring (e..g, [B, C, Z, Y, X])

(same in the class below)

kevinyamauchi · 2024-01-20T15:42:21Z

src/membrain_seg/segmentation/training/surface_dice.py

+    it performs the operation separately for each channel of each batch item.
+    """
+    # Create the Gaussian kernel or load it from the dictionary
+    global gaussian_kernel_dict


Is it necessary to use a global here?

kevinyamauchi · 2024-01-20T15:43:03Z

src/membrain_seg/segmentation/training/surface_dice.py

+        skel = skel + F.relu(delta - skel * delta)
+    return skel
+
+


kevinyamauchi · 2024-01-20T15:43:41Z

tests/membrain_seg/training/test_optim_utils.py

Thanks for updating. It would be nice to increase the test coverage over time, but it's fine for this PR

kevinyamauchi · 2024-01-20T15:45:56Z

src/membrain_seg/segmentation/cli/train_cli.py

@@ -84,6 +87,22 @@ def train_advanced(
            but also severely increases training time.\
                Pass "True" or "False".',
    ),


ah that's a bummer, but I think it's okay. In the future, it might make sense to add support for a glob string or something or directory so that people don't have to write all of the tokens.

kevinyamauchi · 2024-01-20T15:46:44Z

src/membrain_seg/segmentation/training/optim_utils.py

@@ -53,7 +53,7 @@ class IgnoreLabelDiceCELoss(_Loss):
    def __init__(
        self,
        ignore_label: int,


kevinyamauchi · 2024-01-20T15:47:20Z

src/membrain_seg/segmentation/training/optim_utils.py

+        if self.reduction == "mean":
+            combined_loss = combined_loss.mean()
+        elif self.reduction == "sum":
+            combined_loss = combined_loss.sum()


it woudl be nice to add and else case that raises an error so that it doesn't fail silently if people make a typo

LorenzLamm · 2024-01-22T17:45:49Z

This looks good to me! I have some minor comments below. Based on the conversation on zulip with @alisterburt , it sounds like we are in agreement that this loss function should live in membrain-seg. I think you can merge after you address the minor comments. Thanks, @LorenzLamm !

Cool, thanks a lot for your feedback @kevinyamauchi ! Implemented your suggestions and merging now.

LorenzLamm added 18 commits December 31, 2023 14:05

Surface-Dice functionalities

a45d037

Adjust losses to be compatible with Surface-Dice exclusions

df4a376

Adjust training routine and include surface dice loss

79859e0

Add dataset labels to dataloading

4868e98

Pass Surface-Dice arguments to training routine

7bf110b

Update CLI to include advanced options for Surface-Dice

860aa1c

precommit formatting

bd543f5

make list readable by passing argument multiple times

a3c0ce2

remove redundant import

9660790

Compatibility with updated masked_surface_dice function

94a1c81

Add training summary and remove wandb logging

722a5ed

remove reduntant print statements and include ds_labels into for-loop

044a7d6

Implement Gaussian smoothing with torch to compute everything on GPU

84c2b0b

Training summary printing

9b3ae9b

add dataset token to CLI

44bd6a1

Add dataset token to filename

b577c5e

Update warnings

5373977

Fix bug for accuracy masking

b4c9e43

LorenzLamm mentioned this pull request Jan 3, 2024

Potential bug in masked_accuracy. #44

Closed

LorenzLamm commented Jan 3, 2024

View reviewed changes

LorenzLamm added 2 commits January 3, 2024 16:44

Fix Dice reduction to scalar

6d1b0e5

Make test compatible with CombinedLoss

bc8acf0

LorenzLamm commented Jan 3, 2024

View reviewed changes

Fix default path

9f39fc7

LorenzLamm commented Jan 3, 2024

View reviewed changes

kevinyamauchi approved these changes Jan 20, 2024

View reviewed changes

LorenzLamm added 2 commits January 22, 2024 18:31

Raise Error when reduction is not defined

0547611

Add required dimensions to docstrings

346d850

LorenzLamm merged commit 49aa798 into main Jan 22, 2024
5 checks passed

LorenzLamm deleted the SurfaceDice branch January 22, 2024 17:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Surface dice loss #45

Surface dice loss #45

LorenzLamm commented Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

kevinyamauchi Jan 20, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

codecov-commenter commented Jan 3, 2024 •

edited

Loading

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

kevinyamauchi Jan 20, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

kevinyamauchi Jan 20, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

LorenzLamm Jan 3, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi left a comment

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

kevinyamauchi Jan 20, 2024

LorenzLamm commented Jan 22, 2024

		loss_inclusion_tokens.append(surf_dice_tokens)

		scaled_weights = [entry / sum(weights) for entry in weights]

		filtered_seg = F.conv3d(seg, g_kernel, padding=padding, groups=seg.shape[1])
		return filtered_seg

		skel_gt = soft_skel(gt_smooth, iter_=iterations)
		return skel_gt

Surface dice loss #45

Surface dice loss #45

Conversation

LorenzLamm commented Jan 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov-commenter commented Jan 3, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kevinyamauchi left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LorenzLamm commented Jan 22, 2024

codecov-commenter commented Jan 3, 2024 •

edited

Loading