test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class CustomMSE(keras.losses.Loss):\n",
    "#     def __init__(self, regularization_factor=0.1, name=\"custom_mse\"):\n",
    "#         super().__init__(name=name)\n",
    "#         self.regularization_factor = regularization_factor\n",
    "\n",
    "#     def call(self, y_true, y_pred):\n",
    "#         mse = tf.math.reduce_mean(tf.square(y_true - y_pred))\n",
    "#         reg = tf.math.reduce_mean(tf.square(0.5 - y_pred))\n",
    "#         return mse + reg * self.regularization_factor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def cat_cross(y_true, y_pred, weights, epsilon = 1e-9):\n",
    "#     y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "#     y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "#     weights = tf.convert_to_tensor(weights, tf.float32)\n",
    "#     return -tf.reduce_mean(tf.reduce_sum(weights* y_true* tf.math.log(y_pred + epsilon), axis = [-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomWeightedCategoricalCrossentropy(tf.keras.losses.Loss):\n",
    "    def __init__(self, class_weights, epsilon = 1e-9, name = 'wce'):\n",
    "        super().__init__(name=name)\n",
    "        self.class_weights = tf.convert_to_tensor(class_weights, tf.float32)\n",
    "        self.epsilon = epsilon\n",
    "    def call(self, y_true, y_pred):\n",
    "        y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "        y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "        return -tf.reduce_mean(tf.reduce_sum(self.class_weights* y_true* tf.math.log(y_pred + self.epsilon), axis = [-1]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TEST_CASE\n",
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "model = get_unet_128(num_classes = 3)\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = True)\n",
    "\n",
    "\n",
    "l = CustomWeightedCategoricalCrossentropy([1,1,1])\n",
    "print(l(target, prediction))\n",
    "\n",
    "l = tf.keras.losses.CategoricalCrossentropy()\n",
    "print(l(target, prediction))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dice_coeff(y_true, y_pred):\n",
    "    smooth = 1.\n",
    "    y_true_f = K.flatten(y_true)\n",
    "    y_pred_f = K.flatten(y_pred)\n",
    "    intersection = K.sum(y_true_f * y_pred_f)\n",
    "    score = (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)\n",
    "    return score\n",
    "\n",
    "def dice_loss(y_true, y_pred):\n",
    "    y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "    y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "    loss = 1 - dice_coeff(y_true, y_pred)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDiceLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self, smooth = 1.0 ,name = 'dice_loss'):\n",
    "        super().__init__(name = name)\n",
    "        self.smooth = smooth\n",
    "\n",
    "    def compute_dice_coef(self, y_true, y_pred):\n",
    "        intersection = tf.reduce_sum(y_true * y_pred, axis=[1,2,3])\n",
    "        union = tf.reduce_sum(y_true, [1,2,3]) + tf.reduce_sum(y_pred, [1,2,3])\n",
    "        score = (2.0 * intersection + self.smooth) / (union + self.smooth)\n",
    "        return score\n",
    "\n",
    "    def call(self, y_true, y_pred):\n",
    "        y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "        y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "        score  = self.compute_dice_coef(y_true, y_pred)\n",
    "        return 1 - score\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.60010135, shape=(), dtype=float32)\n",
      "\n",
      "tf.Tensor(0.600111, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# TEST_CASE\n",
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "import tensorflow.keras.backend as K\n",
    "model = get_unet_128(num_classes = 3)\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = True)\n",
    "\n",
    "\n",
    "l = CustomDiceLoss()\n",
    "print(l(target, prediction))\n",
    "print()\n",
    "print(dice_loss(target, prediction))\n",
    "\n",
    "# # l = tf.keras.losses.CategoricalCrossentropy()\n",
    "# print(dice_loss(target.astype(np.float32)[...,0][..., np.newaxis], prediction[...,0][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,1][..., np.newaxis], prediction[...,1][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,2][..., np.newaxis], prediction[...,2][..., np.newaxis]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDiceCoefficient(tf.keras.metrics.Metric):\n",
    "    def __init__(self, n_classes, threshold = 0.5, smooth= 1.0, average='macro', name = 'dice_coef', **kwargs):\n",
    "        super().__init__(name = name, **kwargs)\n",
    "        self.n_classes = n_classes\n",
    "        self.smooth = smooth\n",
    "        self.threshold = threshold\n",
    "        self.average = average\n",
    "            \n",
    "        self.sum = self.add_weight(name = 'sum', initializer = 'zeros', shape = (self.n_classes,))\n",
    "        self.total = self.add_weight(name = 'total', initializer = 'zeros')\n",
    "        self.h_iou = self.add_weight(name = 'dice_coef', initializer = 'zeros')\n",
    "    \n",
    "    def update_state(self, y_true, y_pred, sample_weight = None): \n",
    "        correct_values = tf.reduce_sum(tf.cast(tf.argmax(y_pred, axis = 1) == tf.argmax(y_true, axis = 1), \"float32\"))\n",
    "        y_true = tf.cast(tf.cast(y_true, tf.int32), tf.float32)\n",
    "        if self.threshold is not None:\n",
    "            y_pred = tf.cast(tf.cast(y_pred > self.threshold, tf.int32), tf.float32)\n",
    "        else:\n",
    "            y_pred = tf.cast(tf.one_hot( tf.argmax(y_pred, axis = -1), self.n_classes  ), tf.float32)\n",
    "        intersection = tf.reduce_sum(tf.abs(y_true * y_pred), axis=[1,2])\n",
    "        union = tf.reduce_sum(y_true, [1,2]) + tf.reduce_sum(y_pred, [1,2]) - intersection\n",
    "        iou = (intersection + self.smooth) / (union + self.smooth)\n",
    "        \n",
    "        self.sum.assign_add(tf.reduce_sum(iou, axis = 0))\n",
    "        self.total.assign_add(tf.cast(tf.shape(y_true)[0], \"float32\"))\n",
    "            \n",
    "        if self.average == 'macro':\n",
    "            self.h_iou.assign(tf.math.divide_no_nan(tf.reduce_sum(self.weight), tf.reduce_sum(tf.math.divide_no_nan( self.weight, tf.math.divide_no_nan(self.sum, self.total) ))))\n",
    "        elif self.average == 'micro':\n",
    "            self.h_iou.assign( tf.reduce_mean(self.weight * tf.math.divide_no_nan(self.sum, self.total)) )\n",
    "        else:\n",
    "            raise NameError\n",
    "        \n",
    "    def result(self):\n",
    "        return self.h_iou\n",
    "    \n",
    "    def reset_states(self):\n",
    "        self.sum.assign(self.sum * 0.0)\n",
    "        self.total.assign(0.0)\n",
    "        self.h_iou.assign(0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def iou_coef(y_true, y_pred, smooth=1):\n",
    "  intersection = K.sum(K.abs(y_true * y_pred), axis=[1,2,3])\n",
    "  union = K.sum(y_true,[1,2,3])+K.sum(y_pred,[1,2,3])-intersection\n",
    "  iou = (intersection + smooth) / (union + smooth)\n",
    "  return iou"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomIOU(tf.keras.metrics.Metric):\n",
    "    def __init__(self, n_classes, threshold = None, smooth= 1.0, average='macro', weight = None, name = 'weighted_iou', **kwargs):\n",
    "        super().__init__(name = name, **kwargs)\n",
    "        self.n_classes = n_classes\n",
    "        self.smooth = smooth\n",
    "        self.threshold = threshold\n",
    "        self.average = average\n",
    "        if weight is None:\n",
    "            self.weight = tf.ones((n_classes,), tf.float32)\n",
    "        else:\n",
    "            self.weight = tf.convert_to_tensor(weight, tf.float32)\n",
    "            \n",
    "        self.sum = self.add_weight(name = 'sum', initializer = 'zeros', shape = (self.n_classes,))\n",
    "        self.total = self.add_weight(name = 'total', initializer = 'zeros')\n",
    "        self.h_iou = self.add_weight(name = 'h_iou', initializer = 'zeros')\n",
    "    \n",
    "    def update_state(self, y_true, y_pred, sample_weight = None): \n",
    "        y_true = tf.cast(tf.cast(y_true, tf.int32), tf.float32)\n",
    "        if self.threshold is not None:\n",
    "            y_pred = tf.cast(tf.cast(y_pred > self.threshold, tf.int32), tf.float32)\n",
    "        else:\n",
    "            y_pred = tf.cast(tf.one_hot( tf.argmax(y_pred, axis = -1), self.n_classes  ), tf.float32)\n",
    "        intersection = tf.reduce_sum(tf.abs(y_true * y_pred), axis=[1,2])\n",
    "        union = tf.reduce_sum(y_true, [1,2]) + tf.reduce_sum(y_pred, [1,2]) - intersection\n",
    "        iou = (intersection + self.smooth) / (union + self.smooth)\n",
    "        \n",
    "        self.sum.assign_add(tf.reduce_sum(iou, axis = 0))\n",
    "        self.total.assign_add(tf.cast(tf.shape(y_true)[0], \"float32\"))\n",
    "            \n",
    "        if self.average == 'macro':\n",
    "            self.h_iou.assign(tf.math.divide_no_nan(tf.reduce_sum(self.weight), tf.reduce_sum(tf.math.divide_no_nan( self.weight, tf.math.divide_no_nan(self.sum, self.total) ))))\n",
    "        elif self.average == 'micro':\n",
    "            self.h_iou.assign( tf.reduce_mean(self.weight * tf.math.divide_no_nan(self.sum, self.total)) )\n",
    "        else:\n",
    "            raise NameError\n",
    "        \n",
    "    def result(self):\n",
    "        return self.h_iou\n",
    "    \n",
    "    def reset_states(self):\n",
    "        self.sum.assign(self.sum * 0.0)\n",
    "        self.total.assign(0.0)\n",
    "        self.h_iou.assign(0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[0.49306408 0.49815702 0.5038641  0.49650478 0.49640596 0.49932417\n",
      " 0.49223307 0.49229163 0.4938605  0.49797398], shape=(10,), dtype=float32)\n",
      "tf.Tensor(\n",
      "[0.00589825 0.00748396 0.00671221 0.00580411 0.0073432  0.00609088\n",
      " 0.00658778 0.00871226 0.00701585 0.00530376], shape=(10,), dtype=float32)\n",
      "tf.Tensor(\n",
      "[0.00602845 0.00528465 0.00610559 0.00522987 0.00629998 0.00684011\n",
      " 0.0066361  0.00539018 0.00561661 0.00571985], shape=(10,), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "import tensorflow.keras.backend as K\n",
    "model = get_unet_128(num_classes = 3)\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = False)\n",
    "prediction_classes = tf.cast(tf.one_hot( tf.argmax(prediction, axis = -1), 3  ), tf.float32)\n",
    "\n",
    "m = CustomWeightedIOU(3)\n",
    "m.update_state(target, prediction)\n",
    "m.result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tversky_index(y_true, y_pred):\n",
    "    y_true_pos = K.flatten(y_true)\n",
    "    y_pred_pos = K.flatten(y_pred)\n",
    "    true_pos = K.sum(y_true_pos * y_pred_pos)\n",
    "    false_neg = K.sum(y_true_pos * (1 - y_pred_pos))\n",
    "    false_pos = K.sum((1 - y_true_pos) * y_pred_pos)\n",
    "    alpha = 0.7\n",
    "    smooth = 1.0\n",
    "    return (true_pos + smooth) / (true_pos + alpha * false_neg + (\n",
    "                1 - alpha) * false_pos + smooth)\n",
    "\n",
    "def tversky_loss(y_true, y_pred):\n",
    "    y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "    y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "    return 1 - tversky_index(y_true, y_pred)\n",
    "\n",
    "# def focal_tversky(y_true, y_pred):\n",
    "#     pt_1 = tversky_index(y_true, y_pred)\n",
    "#     gamma = 0.75\n",
    "#     return K.pow((1 - pt_1), gamma)\n",
    "\n",
    "# def log_cosh_dice_loss(y_true, y_pred):b\n",
    "#     x = dice_loss(y_true, y_pred)\n",
    "#     return tf.math.log((tf.exp(x) + tf.exp(-x)) / 2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomTverskyLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self, alpha = 0.7, smooth = 1.0 ,name = 'tversky_loss'):\n",
    "        super().__init__(name = name)\n",
    "        self.smooth = smooth\n",
    "        self.alpha = alpha\n",
    "\n",
    "    def compute_tversky_index(self, y_true, y_pred):\n",
    "        true_pos = tf.reduce_sum(y_true * y_pred)\n",
    "        false_neg = tf.reduce_sum(y_true * (1 - y_pred))\n",
    "        false_pos = tf.reduce_sum((1 - y_true) * y_pred)\n",
    "\n",
    "        ti = (true_pos + self.smooth) / (true_pos + self.alpha * false_neg + (1 - self.alpha) * false_pos + self.smooth)\n",
    "        return ti\n",
    "\n",
    "    def call(self, y_true, y_pred):\n",
    "        y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "        y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "        score  = self.compute_tversky_index(y_true, y_pred)\n",
    "        return 1 - score\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.6297133, shape=(), dtype=float32)\n",
      "\n",
      "tf.Tensor(0.6297133, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# TEST_CASE\n",
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "import tensorflow.keras.backend as K\n",
    "model = get_unet_128(num_classes = 3)\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = True)\n",
    "\n",
    "\n",
    "l = CustomTverskyLoss()\n",
    "print(l(target, prediction))\n",
    "print()\n",
    "print(tversky_loss(target, prediction))\n",
    "\n",
    "# # l = tf.keras.losses.CategoricalCrossentropy()\n",
    "# print(dice_loss(target.astype(np.float32)[...,0][..., np.newaxis], prediction[...,0][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,1][..., np.newaxis], prediction[...,1][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,2][..., np.newaxis], prediction[...,2][..., np.newaxis]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "def focal_tversky(y_true, y_pred):\n",
    "    y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "    y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "    pt_1 = tversky_index(y_true, y_pred)\n",
    "    gamma = 0.75\n",
    "    return K.pow((1 - pt_1), gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomFocalTverskyLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self, alpha = 0.7, gamma = 0.75, smooth = 1.0 ,name = 'focal_tversky'):\n",
    "        super().__init__(name = name)\n",
    "        self.smooth = smooth\n",
    "        self.gamma = gamma\n",
    "        self.alpha = alpha\n",
    "\n",
    "    def compute_tversky_index(self, y_true, y_pred):\n",
    "        true_pos = tf.reduce_sum(y_true * y_pred)\n",
    "        false_neg = tf.reduce_sum(y_true * (1 - y_pred))\n",
    "        false_pos = tf.reduce_sum((1 - y_true) * y_pred)\n",
    "\n",
    "        ti = (true_pos + self.smooth) / (true_pos + self.alpha * false_neg + (1 - self.alpha) * false_pos + self.smooth)\n",
    "        return ti\n",
    "\n",
    "    def call(self, y_true, y_pred):\n",
    "        y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "        y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "        score  = self.compute_tversky_index(y_true, y_pred)\n",
    "        return tf.pow(1-score, self.gamma)\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.70692736, shape=(), dtype=float32)\n",
      "\n",
      "tf.Tensor(0.70692736, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# TEST_CASE\n",
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "import tensorflow.keras.backend as K\n",
    "model = get_unet_128(num_classes = 3)\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = True)\n",
    "\n",
    "\n",
    "l = CustomFocalTverskyLoss()\n",
    "print(l(target, prediction))\n",
    "print()\n",
    "print(focal_tversky(target, prediction))\n",
    "\n",
    "# # l = tf.keras.losses.CategoricalCrossentropy()\n",
    "# print(dice_loss(target.astype(np.float32)[...,0][..., np.newaxis], prediction[...,0][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,1][..., np.newaxis], prediction[...,1][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,2][..., np.newaxis], prediction[...,2][..., np.newaxis]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_cosh_dice_loss(y_true, y_pred):\n",
    "    y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "    y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "    x = dice_loss(y_true, y_pred)\n",
    "    return tf.math.log((tf.exp(x) + tf.exp(-x)) / 2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomLogDiceLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self, smooth = 1.0 ,name = 'log_dice_loss'):\n",
    "        super().__init__(name = name)\n",
    "        self.smooth = smooth\n",
    "\n",
    "    def compute_dice_coef(self, y_true, y_pred):\n",
    "        intersection = tf.reduce_sum(y_true * y_pred, axis=[1,2,3])\n",
    "        union = tf.reduce_sum(y_true, [1,2,3]) + tf.reduce_sum(y_pred, [1,2,3])\n",
    "        score = (2.0 * intersection + self.smooth) / (union + self.smooth)\n",
    "        return score\n",
    "\n",
    "    def call(self, y_true, y_pred):\n",
    "        y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "        y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "        score  = self.compute_dice_coef(y_true, y_pred)\n",
    "        return tf.math.log((tf.exp(1-score) + tf.exp(-(1-score))) / 2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(0.17049928, shape=(), dtype=float32)\n",
      "\n",
      "tf.Tensor(0.17050548, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# TEST_CASE\n",
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "import tensorflow.keras.backend as K\n",
    "model = get_unet_128(num_classes = 3)\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = True)\n",
    "\n",
    "\n",
    "l = CustomLogDiceLoss()\n",
    "print(l(target, prediction))\n",
    "print()\n",
    "print(log_cosh_dice_loss(target, prediction))\n",
    "\n",
    "# # l = tf.keras.losses.CategoricalCrossentropy()\n",
    "# print(dice_loss(target.astype(np.float32)[...,0][..., np.newaxis], prediction[...,0][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,1][..., np.newaxis], prediction[...,1][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,2][..., np.newaxis], prediction[...,2][..., np.newaxis]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_dice_coeff(y_true, y_pred, weight):\n",
    "    smooth = 1.\n",
    "    w, m1, m2 = weight * weight, y_true, y_pred\n",
    "    intersection = (m1 * m2)\n",
    "    score = (2. * K.sum(w * intersection) + smooth) / (K.sum(w * m1) + K.sum(w * m2) + smooth)\n",
    "    return score\n",
    "    \n",
    "def weighted_dice_loss(y_true, y_pred):\n",
    "    y_true = K.cast(y_true, 'float32')\n",
    "    y_pred = K.cast(y_pred, 'float32')\n",
    "    # if we want to get same size of output, kernel size must be odd number\n",
    "    if K.int_shape(y_pred)[1] == 128:\n",
    "        kernel_size = 11\n",
    "    elif K.int_shape(y_pred)[1] == 256:\n",
    "        kernel_size = 21\n",
    "    elif K.int_shape(y_pred)[1] == 512:\n",
    "        kernel_size = 21\n",
    "    elif K.int_shape(y_pred)[1] == 1024:\n",
    "        kernel_size = 41\n",
    "    else:\n",
    "        raise ValueError('Unexpected image size')\n",
    "    averaged_mask = K.pool2d(\n",
    "        y_true, pool_size=(kernel_size, kernel_size), strides=(1, 1), padding='same', pool_mode='avg')\n",
    "    border = K.cast(K.greater(averaged_mask, 0.005), 'float32') * K.cast(K.less(averaged_mask, 0.995), 'float32')\n",
    "    weight = K.ones_like(averaged_mask)\n",
    "    w0 = K.sum(weight)\n",
    "    weight += border * 2\n",
    "    w1 = K.sum(weight)\n",
    "    weight *= (w0 / w1)\n",
    "    loss = 1 - weighted_dice_coeff(y_true, y_pred, weight)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def weighted_bce_loss(y_true, y_pred, weight):\n",
    "    # avoiding overflow\n",
    "    epsilon = 1e-7\n",
    "    y_pred = K.clip(y_pred, epsilon, 1. - epsilon)\n",
    "    logit_y_pred = K.log(y_pred / (1. - y_pred))\n",
    "\n",
    "    # https://www.tensorflow.org/api_docs/python/tf/nn/weighted_cross_entropy_with_logits\n",
    "    loss = (1. - y_true) * logit_y_pred + (1. + (weight - 1.) * y_true) * \\\n",
    "                                          (K.log(1. + K.exp(-K.abs(logit_y_pred))) + K.maximum(-logit_y_pred, 0.))\n",
    "    return K.sum(loss) / K.sum(weight)\n",
    "\n",
    "\n",
    "def weighted_bce_dice_loss(y_true, y_pred):\n",
    "    y_true = K.cast(y_true, 'float32')\n",
    "    y_pred = K.cast(y_pred, 'float32')\n",
    "    # if we want to get same size of output, kernel size must be odd number\n",
    "    if K.int_shape(y_pred)[1] == 128:\n",
    "        kernel_size = 11\n",
    "    elif K.int_shape(y_pred)[1] == 256:\n",
    "        kernel_size = 21\n",
    "    elif K.int_shape(y_pred)[1] == 512:\n",
    "        kernel_size = 21\n",
    "    elif K.int_shape(y_pred)[1] == 1024:\n",
    "        kernel_size = 41\n",
    "    else:\n",
    "        raise ValueError('Unexpected image size')\n",
    "    averaged_mask = K.pool2d(\n",
    "        y_true, pool_size=(kernel_size, kernel_size), strides=(1, 1), padding='same', pool_mode='avg')\n",
    "    border = K.cast(K.greater(averaged_mask, 0.005), 'float32') * K.cast(K.less(averaged_mask, 0.995), 'float32')\n",
    "    weight = K.ones_like(averaged_mask)\n",
    "    w0 = K.sum(weight)\n",
    "    weight += border * 2\n",
    "    w1 = K.sum(weight)\n",
    "    weight *= (w0 / w1)\n",
    "    loss = weighted_bce_loss(y_true, y_pred, weight) + (1 - weighted_dice_coeff(y_true, y_pred, weight))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WeightedBCEDiceLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self, pos_weights = None, epsilon = 1e-9, smooth = 1.0, name = 'wbce_dice'):\n",
    "        super().__init__(name = name)\n",
    "        if pos_weights is not None:\n",
    "            self.pos_weights = tf.convert_to_tensor(pos_weights, tf.float32)\n",
    "        else:\n",
    "            self.pos_weights = pos_weights\n",
    "        self.epsilon = epsilon\n",
    "        self.smooth = smooth\n",
    "    def call(self, y_true, y_pred):\n",
    "        y_true = tf.convert_to_tensor(y_true, tf.float32)\n",
    "        y_pred = tf.convert_to_tensor(y_pred, tf.float32)\n",
    "        weight = self.compute_weights(y_true)\n",
    "        if self.pos_weights is not None:\n",
    "            weight = tf.broadcast_to(self.pos_weights, tf.shape(y_true)) + weight \n",
    "        loss = self.weighted_bce_loss(y_true, y_pred, weight) + (1 - self.weighted_dice_coeff(y_true, y_pred, weight))\n",
    "        return loss\n",
    "    def compute_weights(self, y_true):\n",
    "        if tf.shape(y_true)[1] == 128:\n",
    "            kernel_size = 11\n",
    "        elif tf.shape(y_true)[1] == 256:\n",
    "            kernel_size = 21\n",
    "        elif tf.shape(y_true)[1] == 512:\n",
    "            kernel_size = 21\n",
    "        elif tf.shape(y_true)[1] == 1024:\n",
    "            kernel_size = 41\n",
    "        else:\n",
    "            raise ValueError('Unexpected image size')\n",
    "        averaged_mask = K.pool2d(y_true, pool_size=(kernel_size, kernel_size), strides=(1, 1), padding='same', pool_mode='avg')\n",
    "        border = K.cast(K.greater(averaged_mask, 0.005), 'float32') * K.cast(K.less(averaged_mask, 0.995), 'float32')\n",
    "        weight = K.ones_like(averaged_mask)\n",
    "        w0 = K.sum(weight)\n",
    "        weight += border * 2\n",
    "        w1 = K.sum(weight)\n",
    "        weight *= (w0 / w1)\n",
    "        return tf.stop_gradient(weight)\n",
    "\n",
    "    def weighted_bce_loss(self, y_true, y_pred, weight):\n",
    "        y_pred = K.clip(y_pred, self.epsilon, 1. - self.epsilon)\n",
    "        logit_y_pred = K.log(y_pred / (1. - y_pred))\n",
    "\n",
    "        # https://www.tensorflow.org/api_docs/python/tf/nn/weighted_cross_entropy_with_logits\n",
    "        loss = (1. - y_true) * logit_y_pred + (1. + (weight - 1.) * y_true) * \\\n",
    "                                            (K.log(1. + K.exp(-K.abs(logit_y_pred))) + K.maximum(-logit_y_pred, 0.))\n",
    "        return K.sum(loss) / K.sum(weight)\n",
    "\n",
    "    def weighted_dice_loss(self,y_true, y_pred):\n",
    "        weight = self.compute_weights(y_true)\n",
    "        loss = 1 - weighted_dice_coeff(y_true, y_pred, weight)\n",
    "        return loss\n",
    "\n",
    "    def weighted_dice_coeff(self, y_true, y_pred, weight):\n",
    "        w, m1, m2 = weight * weight, y_true, y_pred\n",
    "        intersection = (m1 * m2)\n",
    "        score = (2. * K.sum(w * intersection) + self.smooth) / (K.sum(w * m1) + K.sum(w * m2) + self.smooth)\n",
    "        return score\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(1.3155478, shape=(), dtype=float32)\n",
      "\n",
      "tf.Tensor(1.3155478, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# TEST_CASE\n",
    "from unet import get_unet_128\n",
    "import numpy as np\n",
    "import tensorflow.keras.backend as K\n",
    "model = get_unet_128(num_classes = 3, final_activation='sigmoid')\n",
    "x = np.random.uniform(0,1, (10,128,128,3))\n",
    "target = np.random.randint(0,2, (10,128,128,3))\n",
    "prediction = model(x, training = True)\n",
    "\n",
    "\n",
    "l = WeightedBCEDiceLoss()\n",
    "print(l(target, prediction))\n",
    "print()\n",
    "print(weighted_bce_dice_loss(target, prediction))\n",
    "\n",
    "# # l = tf.keras.losses.CategoricalCrossentropy()\n",
    "# print(dice_loss(target.astype(np.float32)[...,0][..., np.newaxis], prediction[...,0][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,1][..., np.newaxis], prediction[...,1][..., np.newaxis]))\n",
    "# print(dice_loss(target.astype(np.float32)[...,2][..., np.newaxis], prediction[...,2][..., np.newaxis]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10, 128, 128, 3), dtype=float32, numpy=\n",
       "array([[[[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]]],\n",
       "\n",
       "\n",
       "       [[[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]]],\n",
       "\n",
       "\n",
       "       [[[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]]],\n",
       "\n",
       "\n",
       "       ...,\n",
       "\n",
       "\n",
       "       [[[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]]],\n",
       "\n",
       "\n",
       "       [[[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]]],\n",
       "\n",
       "\n",
       "       [[[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]],\n",
       "\n",
       "        [[1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         ...,\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.],\n",
       "         [1., 1., 1.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.compute_weights(target.astype(np.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}