deltaE2000 loss function

deltaE2000_loss.py

# Code adapted from skimage
# https://github.com/riaanvddool/scikits-image/blob/master/skimage/color/delta_e.py#L123

import numpy as np
import tensorflow as tf

def tf_deg2rad(deg):
    return (deg / 180) * np.pi

def tf_rad2deg(rad):
    return (rad / np.pi) * 180

def tf_cart2polar_2pi(x, y):
    """convert cartesian coordiantes to polar (uses non-standard theta range!)
    NON-STANDARD RANGE! Maps to (0, 2*pi) rather than usual (-pi, +pi)
    """
    r, t = tf.math.sqrt(x ** 2 + y ** 2), tf.math.atan2(y, x)
    t += tf.where(t < 0., 2 * np.pi, 0)
    return r, t

def tf_deltaE2000(lab1, lab2, kL=1, kC=1, kH=1):
  L1 = lab1[:, 0]
  a1 = lab1[:, 1]
  b1 = lab1[:, 2]
  L2 = lab2[:, 0]
  a2 = lab2[:, 1]
  b2 = lab2[:, 2]
  
  # distort `a` based on average chroma
  # then convert to lch coordines from distorted `a`
  # all subsequence calculations are in the new coordiantes
  # (often denoted "prime" in the literature)
  Cbar = 0.5 * (tf.sqrt(a1 **2 + b1 **2) + tf.sqrt(a2 ** 2 + b2 ** 2))
  c7 = Cbar ** 7
  G = 0.5 * (1 - tf.sqrt(c7 / (c7 + 25 ** 7)))
  scale = 1 + G
  C1, h1 = tf_cart2polar_2pi(a1 * scale, b1)
  C2, h2 = tf_cart2polar_2pi(a2 * scale, b2)
  # # recall that c, h are polar coordiantes.  c==r, h==theta

  # # cide2000 has four terms to delta_e:
  # # 1) Luminance term
  # # 2) Hue term
  # # 3) Chroma term
  # # 4) hue Rotation term

  # # lightness term
  Lbar = 0.5 * (L1 + L2)
  tmp = (Lbar - 50) ** 2
  SL = 1 + 0.015 * tmp / tf.math.sqrt(20 + tmp)
  L_term = (L2 - L1) / (kL * SL)

  # # chroma term
  Cbar = 0.5 * (C1 + C2)  # new coordiantes
  SC = 1 + 0.045 * Cbar
  C_term = (C2 - C1) / (kC * SC)

  # # hue term
  h_diff = h2 - h1
  h_sum = h1 + h2
  CC = C1 * C2

  dH = tf.identity(h_diff)
  dH = tf.where(h_diff > np.pi, dH - 2 * np.pi, dH)
  dH = tf.where(h_diff < -np.pi, dH + 2 * np.pi, dH)
  dH = tf.where(CC < 0., 0., dH)
  dH_term = 2 * tf.math.sqrt(CC) * tf.math.sin(dH / 2)
  
  Hbar = tf.identity(h_sum)
  mask = tf.math.logical_and(CC != 0., tf.abs(h_diff) > np.pi)
  mask2 = tf.math.logical_and(mask, h_sum >= 2 * np.pi)
  Hbar = tf.where(mask2, Hbar - 2 * np.pi, Hbar)
  Hbar = tf.where(CC == 0., Hbar * 2, Hbar)
  Hbar *= 0.5

  T = (1 -
      0.17 * tf.math.cos(Hbar - tf_deg2rad(30)) +
      0.24 * tf.math.cos(2 * Hbar) +
      0.32 * tf.math.cos(3 * Hbar + tf_deg2rad(6)) -
      0.20 * tf.math.cos(4 * Hbar - tf_deg2rad(63))
      )
  SH = 1 + 0.015 * Cbar * T

  H_term = dH_term / (kH * SH)

  # # hue rotation
  c7 = Cbar ** 7
  Rc = 2 * tf.math.sqrt(c7 / (c7 + 25 ** 7))
  dtheta = tf_deg2rad(30) * tf.math.exp(-((tf_rad2deg(Hbar) - 275) / 25) ** 2)
  R_term = -tf.math.sin(2 * dtheta) * Rc * C_term * H_term
  
  # # put it all together
  dE2 = L_term ** 2
  dE2 += C_term ** 2
  dE2 += H_term ** 2
  dE2 += R_term
  return tf.math.sqrt(dE2)