segmentation loss

From original article:
https://lars76.github.io/neural-networks/object-detection/losses-for-segmentation/

IoU loss https://www.cs.umanitoba.ca/~ywang/papers/isvc16.pdf

Focal loss

Focal loss (FL) [2] tries to down-weight the contribution of easy examples so that the CNN focuses more on hard examples.

FL can be defined as follows:

$$\text{FL}(p,\hat{p})=-\left(\alpha \right(1-\hat{p}{)}^{\gamma }p\log \left(\hat{p}\right)+(1-\alpha ){\hat{p}}^{\gamma }(1-p)\log (1-\hat{p}))$$

When $\gamma =0$, we obtain BCE.

This time we cannot use weighted_cross_entropy_with_logits to implement FL in Keras. We will derive instead our own focal_loss_with_logits function.

$$\begin{array}{cc}& =\alpha (1-\hat{p}{)}^{\gamma }p\log (1+e^{-x})-(1-\alpha ){\hat{p}}^{\gamma }(1-p)\log \left(\frac{e^{-x}}{1+e^{-x}}\right)\\ & =\alpha (1-\hat{p}{)}^{\gamma }p\log (1+e^{-x})-(1-\alpha ){\hat{p}}^{\gamma }(1-p)(-x-\log (1+e^{-x}))\\ & =\alpha (1-\hat{p}{)}^{\gamma }p\log (1+e^{-x})+(1-\alpha ){\hat{p}}^{\gamma }(1-p)(x+\log (1+e^{-x}))\\ & =\log (1+e^{-x})\left(\alpha \right(1-\hat{p}{)}^{\gamma }p+(1-\alpha ){\hat{p}}^{\gamma }(1-p))+x(1-\alpha ){\hat{p}}^{\gamma }(1-p)\\ & =\log \left(e^{-x}\right(1+e^x\left)\right)\left(\alpha \right(1-\hat{p}{)}^{\gamma }p+(1-\alpha ){\hat{p}}^{\gamma }(1-p))+x(1-\alpha ){\hat{p}}^{\gamma }(1-p)\\ & =(\log (1+e^x)-x)\left(\alpha \right(1-\hat{p}{)}^{\gamma }p+(1-\alpha ){\hat{p}}^{\gamma }(1-p))+x(1-\alpha ){\hat{p}}^{\gamma }(1-p)\\ & =(\log (1+e^{-|x|})+max(-x,0\left)\right)\left(\alpha \right(1-\hat{p}{)}^{\gamma }p+(1-\alpha ){\hat{p}}^{\gamma }(1-p))+x(1-\alpha ){\hat{p}}^{\gamma }(1-p)\end{array}$$

And the implementation is then:

def focal_loss(alpha=0.25, gamma=2):
  def focal_loss_with_logits(logits, targets, alpha, gamma, y_pred):
    weight_a = alpha * (1 - y_pred) ** gamma * targets
    weight_b = (1 - alpha) * y_pred ** gamma * (1 - targets)
    
    return (tf.log1p(tf.exp(-tf.abs(logits))) + tf.nn.relu(-logits)) * (weight_a + weight_b) + logits * weight_b 

  def loss(y_true, y_pred):
    y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon())
    logits = tf.log(y_pred / (1 - y_pred))

    loss = focal_loss_with_logits(logits=logits, targets=y_true, alpha=alpha, gamma=gamma, y_pred=y_pred)

    return tf.reduce_mean(loss)

  return loss

Overlap measures

Dice Loss / F1 score

The Dice coefficient is similar to the Jaccard Index (Intersection over Union, IoU):

$$\text{DC}=\frac{2TP}{2TP+FP+FN}=\frac{2|X\cap Y|}{|X|+|Y|}$$$$\text{IoU}=\frac{TP}{TP+FP+FN}=\frac{|X\cap Y|}{|X|+|Y|-|X\cap Y|}$$

where TP are the true positives, FP false positives and FN false negatives. We can see that $\text{DC}\ge \text{IoU}$.

The dice coefficient can also be defined as a loss function:

$$\text{DL}(p,\hat{p})=\frac{2⟨p,\hat{p}⟩}{{\parallel p\parallel }_2^2+{\parallel \hat{p}\parallel }_2^2}$$

where $p\in \{0,1{\}}^n$ and $0\le \hat{p}\le 1$.

def dice_loss(y_true, y_pred):
  numerator = 2 * tf.reduce_sum(y_true * y_pred)
  # some implementations don't square y_pred
  denominator = tf.reduce_sum(y_true + tf.square(y_pred))

  return numerator / (denominator + tf.keras.backend.epsilon())

Since $p$ is either $1$ or $0$, the numerator will always be one times the predicted probability of the foreground pixel (1). Hence, when $p$ is a background pixel (0), the numerator will be 0.

Tversky loss

Tversky loss (TL) is a generalization of Dice loss. TL adds a weight to FP and FN.

$$\text{DL}(p,\hat{p})=\frac{⟨p,\hat{p}⟩}{⟨p,\hat{p}⟩+\beta ⟨1-p,\hat{p}⟩+(1-\beta )⟨p,1-\hat{p}⟩}$$

Let $\beta =\frac{1}{2}$. Then

$$\begin{array}{cc}& =\frac{2⟨p,\hat{p}⟩}{2⟨p,\hat{p}⟩+⟨1-p,\hat{p}⟩+⟨p,1-\hat{p}⟩}\\ & =\frac{2⟨p,\hat{p}⟩}{⟨1,\hat{p}⟩+⟨1,p⟩}\end{array}$$

which is just Dice loss. In the paper [4], the authors square the predicted probability in the denominator, but e.g. the paper [5]keeps the term as it is.

def tversky_loss(beta):
  def loss(y_true, y_pred):
    numerator = tf.reduce_sum(y_true * y_pred)
    denominator = y_true * y_pred + beta * (1 - y_true) * y_pred + (1 - beta) * y_true * (1 - y_pred)

    return numerator / (tf.reduce_sum(denominator) + tf.keras.backend.epsilon())

  return loss

Lovász-Softmax

DL and TL simply relax the hard constraint $\hat{p}\in \{0,1{\}}^n$ in order to have a function on the domain $[0,1]$. The paper [6] derives instead a surrogate loss function.

An implementation of Lovász-Softmax can be found on github. Note that this loss requires the identity activation in the last layer. A negative value means class A and a positive value means class B.

In Keras the loss function can be used as follows:

def lovasz_softmax(y_true, y_pred):
  return lovasz_hinge(labels=y_true, logits=y_pred)

model.compile(loss=lovasz_softmax, optimizer=optimizer, metrics=[pixel_iou])

References

[1] S. Xie and Z. Tu. Holistically-Nested Edge Detection, 2015.

[2] T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Dollar. Focal Loss for Dense Object Detection, 2017.

[3] O. Ronneberger, P. Fischer, and T. Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation, 2015.

[4] F. Milletari, N. Navab, and S.-A. Ahmadi. V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation, 2016.

[5] S. S. M. Salehi, D. Erdogmus, and A. Gholipour. Tversky loss function for image segmentation using 3D fully convolutional deep networks, 2017.

[6] M. Berman, A. R. Triki, M. B. Blaschko. The Lovász-Softmax loss: A tractable surrogate for the optimization of the intersection-over-union measure in neural networks, 2018.