Surface loss in keras-tensorflow

[AmericaBG]

Hi!
Could you help me with the implementation of surface loss function in keras and tensorflow?
Thank you very much in advance!

My best wishes.

[HKervadec]

Hey,

This should actually be quite easy. The most difficult part is actually generating the distance maps from the ground truth, and fortunately this is done in NumPy. I don't know much about Keras or Tensorflow, but I assume you can either do that asynchronously in some data loader, or you could do that offline once.

Assuming you have a class-encoded ground truth as an array of shape wh, where each pixel is the class number, the first step would be to move to some one-hot encoding. A function like this should do:

def numpy_class2one_hot(seg: np.ndarray, C: int) -> np.ndarray:
    assert set(np.unique(seg)).issubset(list(range(C)))

    w, h = seg.shape  # type: Tuple[int, int, int]

    res = np.stack([seg == c for c in range(C)], axis=0).astype(np.int32)
    assert res.shape == (C, w, h)
    assert np.all(res.sum(axis=0) == 1)

    return res

This will output a one-hot encoded image of shape cwh.

Then, you could plug the results directly into the one_hot2dist function defined in utils.py:

def one_hot2dist(seg: np.ndarray) -> np.ndarray:
    assert one_hot(torch.Tensor(seg), axis=0)
    C: int = len(seg)

    res = np.zeros_like(seg)
    for c in range(C):
        posmask = seg[c].astype(np.bool)

        if posmask.any():
            negmask = ~posmask
            res[c] = distance(negmask) * negmask - (distance(posmask) - 1) * posmask
    return res

This will compute one distance map per class, for an output shape cwh.

Then, you can load that in your favorite framework, and multiply it element-wise to your network softmax probabilities (which should be of shape cwh as well). I personally then average the results, but this is a matter of taste, summing them would be fine also (it merely affects the learning rate).

Hope that helps, let me know if you need any help.

[che85]

@HKervadec FIrst of all, a very nice paper! Do you have any test data and expected results when running the loss function? I wanted to try your loss on my data (different project template) but I am not sure if it's actually spitting out the correct loss. Do you have any recommendations on how to test this?

Thank you,
Christian

[Eason270]

Could you provide the implement of surface_loss in Keras?

[HKervadec]

@HKervadec FIrst of all, a very nice paper! Do you have any test data and expected results when running the loss function? I wanted to try your loss on my data (different project template) but I am not sure if it's actually spitting out the correct loss. Do you have any recommendations on how to test this?

Hey sorry for the slow reply. I do not have much test data/values, but I had this testing script that I made a while back that you could use as a basis

#!/usr/bin/env python3.6

import unittest

import torch
import numpy as np

import utils

class TestDistMap(unittest.TestCase):
    def test_closure(self):
        a = np.zeros((1, 256, 256))
        a[:, 50:60, :] = 1

        o = utils.class2one_hot(torch.Tensor(a).type(torch.float32), C=2).numpy()
        res = utils.one_hot2dist(o[0])
        self.assertEqual(res.shape, (2, 256, 256))

        neg = (res <= 0) * res

        self.assertEqual(neg.sum(), (o * res).sum())

    def test_full_coverage(self):
        a = np.zeros((1, 256, 256))
        a[:, 50:60, :] = 1

        o = utils.class2one_hot(torch.Tensor(a).type(torch.float32), C=2).numpy()
        res = utils.one_hot2dist(o[0])
        self.assertEqual(res.shape, (2, 256, 256))

        self.assertEqual((res[1] <= 0).sum(), a.sum())
        self.assertEqual((res[1] > 0).sum(), (1 - a).sum())

    def test_empty(self):
        a = np.zeros((1, 256, 256))

        o = utils.class2one_hot(torch.Tensor(a).type(torch.float32), C=2).numpy()
        res = utils.one_hot2dist(o[0])
        self.assertEqual(res.shape, (2, 256, 256))

        self.assertEqual(res[1].sum(), 0)
        self.assertEqual((res[0] <= 0).sum(), a.size)

    def test_max_dist(self):
        """
        The max dist for a box should be at the midle of the object, +-1
        """
        a = np.zeros((1, 256, 256))
        a[:, 1:254, 1:254] = 1

        o = utils.class2one_hot(torch.Tensor(a).type(torch.float32), C=2).numpy()
        res = utils.one_hot2dist(o[0])
        self.assertEqual(res.shape, (2, 256, 256))

        self.assertEqual(res[0].max(), 127)
        self.assertEqual(np.unravel_index(res[0].argmax(), (256, 256)), (127, 127))

        self.assertEqual(res[1].min(), -126)
        self.assertEqual(np.unravel_index(res[1].argmin(), (256, 256)), (127, 127))

    def test_border(self):
        """
        Make sure the border inside the object is 0 in the distance map
        """

        for l in range(3, 5):
            a = np.zeros((1, 25, 25))
            a[:, 3:3 + l, 3:3 + l] = 1

            o = utils.class2one_hot(torch.Tensor(a).type(torch.float32), C=2).numpy()
            res = utils.one_hot2dist(o[0])
            self.assertEqual(res.shape, (2, 25, 25))

            border = (res[1] == 0)

            self.assertEqual(border.sum(), 4 * (l - 1))


if __name__ == "__main__":
    unittest.main()

utils refers to utils.py in the main repository (or maybe a more recent version of it, but no big update)

Could you provide the implement of surface_loss in Keras?

I'm sorry, I do not have the time to do that, not the means to test it afterwards. But I think I gave enough informations in #14 (comment) to do it.

I would be willing to publish in this repository a Keras implementation that someone submit as a PR.

[warmestwind]

@Eason270 @SouthAmericaB
For Tensorflow or Keras, just use tf.py_func to wrap numpy_class2one_hot and one_hot2dist then use them as a TensorFlow OP.

[akamojo]

Hi,
@HKervadec, thanks to the surface loss we managed to improve our results of i.a. hippocampus segmentation task :)
This is how we implemented surface loss in keras, together with @marcinkaczor:

from keras import backend as K
import numpy as np
import tensorflow as tf
from scipy.ndimage import distance_transform_edt as distance


def calc_dist_map(seg):
    res = np.zeros_like(seg)
    posmask = seg.astype(np.bool)

    if posmask.any():
        negmask = ~posmask
        res = distance(negmask) * negmask - (distance(posmask) - 1) * posmask

    return res

def calc_dist_map_batch(y_true):
    y_true_numpy = y_true.numpy()
    return np.array([calc_dist_map(y)
                     for y in y_true_numpy]).astype(np.float32)

def surface_loss(y_true, y_pred):
    y_true_dist_map = tf.py_function(func=calc_dist_map_batch,
                                     inp=[y_true],
                                     Tout=tf.float32)
    multipled = y_pred * y_true_dist_map
    return K.mean(multipled)

However, this is an implementation only for the problem of binary classification (foreground and background).

[HKervadec]

This looks great.

However (I am not familiar much with Keras and Tensorflow), the distance map isn't precomputed ; meaning it has to be done when computing the loss, which I assume slow things down.

Is there some scheduling / compiling process that optimize the order of the operation ? From what I recall Tensorflow build its graph then runs into a session ; a clever compiler would probably reorder those parts to avoid waiting for it.

[AmericaBG]

Hi,
@HKervadec, thanks to the surface loss we managed to improve our results of i.a. hippocampus segmentation task :)
This is how we implemented surface loss in keras, together with @marcinkaczor:

from keras import backend as K
import numpy as np
import tensorflow as tf
from scipy.ndimage import distance_transform_edt as distance


def calc_dist_map(seg):
    res = np.zeros_like(seg)
    posmask = seg.astype(np.bool)

    if posmask.any():
        negmask = ~posmask
        res = distance(negmask) * negmask - (distance(posmask) - 1) * posmask

    return res

def calc_dist_map_batch(y_true):
    y_true_numpy = y_true.numpy()
    return np.array([calc_dist_map(y)
                     for y in y_true_numpy]).astype(np.float32)

def surface_loss(y_true, y_pred):
    y_true_dist_map = tf.py_function(func=calc_dist_map_batch,
                                     inp=[y_true],
                                     Tout=tf.float32)
    multipled = y_pred * y_true_dist_map
    return K.mean(multipled)

However, this is an implementation only for the problem of binary classification (foreground and background).

Hi! Thank you very much for your code!
I've tried to use it as loss function in a binary classification task but I haven't got good results (my accuracy index tends to 0 instead of 1.

I use your surface_loss function when model.compile(....loss=[surface_loss])

Have you any idea what could be happening?

Thank you very much in advance!

[marcinkaczor]

@HKervadec, we considered computing distance maps before network training, but we noticed that it might be hard to associate a label (image mask) with a proper distance map during training (dataset is shuffled). So that we are doing it on the fly.
@SouthAmericaB, using a surface loss alone doesn't give good results also with us. We combined a surface loss with a generalised dice loss as described in the paper. We used a callback to schedule alpha value for each epoch:

class AlphaScheduler(Callback):
  def init(self, alpha, update_fn):
    self.alpha = alpha
    self.update_fn = update_fn
  def on_epoch_end(self, epoch, logs=None):
    updated_alpha = self.update_fn(K.get_value(self.alpha))

alpha = K.variable(1, dtype='float32')

def update_alpha(value):
  return np.clip(value - 0.01, 0.01, 1)

history = model.fit_generator(
  ...,
  callbacks=AlphaScheduler(alpha, update_alpha)
)

[HKervadec]

we noticed that it might be hard to associate a label (image mask) with a proper distance map during training (dataset is shuffled).

I am not quite sure I understand. How come it is harder to keep track of a tuple (input, dist_transform(label)) instead of (input. label) ? I mean, you already need to maintain the association for standard losses, so I don't get why this is different when you transform the label further. Is there something in the Keras/Tensorflow paradigm that makes it more difficult ?

In any case, thanks for all the contributions ! The scheduler bit is very useful as well ; I will integrate it in the pull request as well.

[marcinkaczor]

@HKervadec, thank you for sharing your improvement idea. We would have to take a closer look at whether such an implementation would be possible in our case.

[AmericaBG]

@HKervadec, we considered computing distance maps before network training, but we noticed that it might be hard to associate a label (image mask) with a proper distance map during training (dataset is shuffled). So that we are doing it on the fly.
@SouthAmericaB, using a surface loss alone doesn't give good results also with us. We combined a surface loss with a generalised dice loss as described in the paper. We used a callback to schedule alpha value for each epoch:

class AlphaScheduler(Callback):
  def init(self, alpha, update_fn):
    self.alpha = alpha
    self.update_fn = update_fn
  def on_epoch_end(self, epoch, logs=None):
    updated_alpha = self.update_fn(K.get_value(self.alpha))

alpha = K.variable(1, dtype='float32')

def update_alpha(value):
  return np.clip(value - 0.01, 0.01, 1)

history = model.fit_generator(
  ...,
  callbacks=AlphaScheduler(alpha, update_alpha)
)

Thank you very much about your explanation!
But now I don't understand what the alpha paramether means

@HKervadec, we considered computing distance maps before network training, but we noticed that it might be hard to associate a label (image mask) with a proper distance map during training (dataset is shuffled). So that we are doing it on the fly.
@SouthAmericaB, using a surface loss alone doesn't give good results also with us. We combined a surface loss with a generalised dice loss as described in the paper. We used a callback to schedule alpha value for each epoch:

class AlphaScheduler(Callback):
  def init(self, alpha, update_fn):
    self.alpha = alpha
    self.update_fn = update_fn
  def on_epoch_end(self, epoch, logs=None):
    updated_alpha = self.update_fn(K.get_value(self.alpha))

alpha = K.variable(1, dtype='float32')

def update_alpha(value):
  return np.clip(value - 0.01, 0.01, 1)

history = model.fit_generator(
  ...,
  callbacks=AlphaScheduler(alpha, update_alpha)
)

Oh! That's true! Thank you very much!
One last thing, could you give me the generalised dice loss function in keras-tensorflow??

You will be doing me a great service!

[HKervadec]

One last thing, could you give me the generalised dice loss function in keras-tensorflow??

You are not limited to GDL for the regional loss ; any other can work (cross-entropy and its derivative, dice loss and its derivatives). The best one will depend on your specific application, but you can already try with others.