Added loss functions and metrices

.gitignore

@@ -105,4 +105,8 @@ ENV/
 .ropeproject
 
 # VSCode settings
-.vscode/
+.vscode/
+.idea/image-segmentation-keras.iml
+.idea/inspectionProfiles/profiles_settings.xml
+.idea/modules.xml
+.idea/inspectionProfiles/Project_Default.xml

README.md

@@ -4,7 +4,7 @@
 [![Downloads](https://pepy.tech/badge/keras-segmentation)](https://pepy.tech/project/keras-segmentation)
 [![Build Status](https://travis-ci.org/divamgupta/image-segmentation-keras.png)](https://travis-ci.org/divamgupta/image-segmentation-keras)
 [![MIT license](https://img.shields.io/badge/License-MIT-blue.svg)](http://perso.crans.org/besson/LICENSE.html)
-[![Twitter](https://img.shields.io/twitter/url.svg?label=Follow%20%40divamgupta&style=social&url=https%3A%2F%2Ftwitter.com%2Fdivamgupta)](https://twitter.com/divamgupta)
+[![Twitter](https://img.shields.io/twitter/url.svg?label=Follow%20%40sushantag9&style=social&url=https%3A%2F%2Ftwitter.com%2Fdivamgupta)](https://twitter.com/sushantag9)
 
 
 
@@ -170,10 +170,14 @@ from keras_segmentation.models.unet import vgg_unet
 
 model = vgg_unet(n_classes=51 ,  input_height=416, input_width=608  )
 
+def custom_fun():
+    #some loss function
+    pass
+
 model.train(
     train_images =  "dataset1/images_prepped_train/",
     train_annotations = "dataset1/annotations_prepped_train/",
-    checkpoints_path = "/tmp/vgg_unet_1" , epochs=5
+    checkpoints_path = "/tmp/vgg_unet_1" , epochs=5, loss_fun = custom_fun() #dice or soft_dice, default = 'categorical_crossentropy'
 )
 
 out = model.predict_segmentation(
@@ -186,9 +190,14 @@ plt.imshow(out)
 
 # evaluating the model 
 print(model.evaluate_segmentation( inp_images_dir="dataset1/images_prepped_test/"  , annotations_dir="dataset1/annotations_prepped_test/" ) )
-
 ```
+### Custom loss function
 
+#### Available loss function:
+
+- Dice loss (usage: `loss_fun = dice`)
+- Soft Dice Loss (usage: `loss_fun = soft_dice`)
+- and others available in [Losses](https://keras.io/api/losses/)
 
 ## Usage via command line
 You can also use the tool just using command line
@@ -229,6 +238,7 @@ python -m keras_segmentation train \
  --input_height=320 \
  --input_width=640 \
  --model_name="vgg_unet"
+ --loss_fun= "dice" or "soft_dice" #default = 'categorical_crossentropy'
 ```
 
 Choose model_name from the table above

keras_segmentation/metrics.py

@@ -1,4 +1,5 @@
 import numpy as np
+from tensorflow.keras import backend as K
 
 EPS = 1e-12
 
@@ -11,3 +12,23 @@ def get_iou(gt, pr, n_classes):
         iou = float(intersection)/(union + EPS)
         class_wise[cl] = iou
     return class_wise
+
+def dice_coef(y_true, y_pred):
+
+    y_true_f = K.flatten(y_true)
+    y_pred_f = K.flatten(y_pred)
+    intersection = K.sum(y_true_f * y_pred_f)
+    return  (2. * intersection + EPS) / (K.sum(y_true_f) + K.sum(y_pred_f) + EPS)
+
+# def dice_coef(y_true, y_pre, smooth=1):
+#     # Dice Coefficient is 2 * the Area of Overlap divided by the total number of pixels in both images
+#     y_true = np.asarray(y_true, 'bool')
+#     y_pre = np.asarray(y_pre, 'bool')
+#     inter = np.sum(y_true * y_pre)
+#     uni = np.sum(y_true) + np.sum(y_pre)
+#
+#     return (2 * inter + smooth) / (uni + smooth)
+
+
+def jacard_index(dice):
+    return dice / (2 - dice)

keras_segmentation/model_compression.py

@@ -1,4 +1,4 @@
-import keras
+import tensorflow.keras as keras
 import tensorflow as tf
 
 from tqdm import tqdm 
@@ -11,7 +11,7 @@
 from .data_utils.data_loader import image_segmentation_generator
 from .train import CheckpointsCallback 
 
-from keras.models import Model 
+from tensorflow.keras.models import Model
 
 
 

keras_segmentation/models/_pspnet_2.py

@@ -2,11 +2,11 @@
 
 from math import ceil
 from sys import exit
-from keras import layers
-from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D
-from keras.layers import BatchNormalization, Activation,\
+from tensorflow.keras import layers
+from tensorflow.keras.layers import Conv2D, MaxPooling2D, AveragePooling2D
+from tensorflow.keras.layers import BatchNormalization, Activation,\
     Input, Dropout, ZeroPadding2D
-from keras.layers.merge import Concatenate, Add
+from tensorflow.keras.layers.merge import Concatenate, Add
 import tensorflow as tf
 
 from .config import IMAGE_ORDERING

keras_segmentation/models/basic_models.py

@@ -1,6 +1,6 @@
-from keras.models import *
-from keras.layers import *
-import keras.backend as K
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
+import tensorflow.keras.backend as K
 
 from .config import IMAGE_ORDERING
 

keras_segmentation/models/fcn.py

@@ -1,5 +1,5 @@
-from keras.models import *
-from keras.layers import *
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
 
 from .config import IMAGE_ORDERING
 from .model_utils import get_segmentation_model

keras_segmentation/models/mobilenet.py

@@ -1,7 +1,7 @@
-from keras.models import *
-from keras.layers import *
-import keras.backend as K
-import keras
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
+import tensorflow.keras.backend as K
+import tensorflow.keras as keras
 
 from .config import IMAGE_ORDERING
 

keras_segmentation/models/model_utils.py

@@ -1,8 +1,8 @@
 from types import MethodType
 
-from keras.models import *
-from keras.layers import *
-import keras.backend as K
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
+import tensorflow.keras.backend as K
 from tqdm import tqdm
 
 from .config import IMAGE_ORDERING

keras_segmentation/models/pspnet.py

@@ -1,8 +1,8 @@
 import numpy as np
-import keras
-from keras.models import *
-from keras.layers import *
-import keras.backend as K
+import tensorflow.keras as keras
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
+import tensorflow.keras.backend as K
 
 from .config import IMAGE_ORDERING
 from .model_utils import get_segmentation_model, resize_image

keras_segmentation/models/resnet50.py

@@ -1,7 +1,7 @@
-import keras
-from keras.models import *
-from keras.layers import *
-from keras import layers
+import tensorflow.keras as keras
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
+from tensorflow.keras import layers
 
 # Source:
 # https://github.com/fchollet/deep-learning-models/blob/master/resnet50.py

keras_segmentation/models/segnet.py

@@ -1,5 +1,5 @@
-from keras.models import *
-from keras.layers import *
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
 
 from .config import IMAGE_ORDERING
 from .model_utils import get_segmentation_model

keras_segmentation/models/unet.py

@@ -1,5 +1,5 @@
-from keras.models import *
-from keras.layers import *
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
 
 from .config import IMAGE_ORDERING
 from .model_utils import get_segmentation_model

keras_segmentation/models/vgg16.py

@@ -1,6 +1,6 @@
-import keras
-from keras.models import *
-from keras.layers import *
+import tensorflow.keras as keras
+from tensorflow.keras.models import *
+from tensorflow.keras.layers import *
 
 from .config import IMAGE_ORDERING
 

keras_segmentation/pretrained.py

@@ -1,4 +1,5 @@
-import keras
+# import keras
+import tensorflow.keras as keras
 
 from .models.all_models import model_from_name
 

keras_segmentation/train.py

@@ -7,9 +7,28 @@
 from keras.callbacks import Callback
 from tensorflow.keras.callbacks import ModelCheckpoint
 import tensorflow as tf
+from tensorflow.keras import backend as K
 import glob
 import sys
 
+
+def dice_coef(y_true, y_pred):
+
+    y_true_f = K.flatten(y_true)
+    y_pred_f = K.flatten(y_pred)
+    intersection = K.sum(y_true_f * y_pred_f)
+    return (2. * intersection + K.epsilon()) / (K.sum(y_true_f) + K.sum(y_pred_f) + K.epsilon())
+
+def dice_coef_loss(y_true, y_pred):
+    return 1-dice_coef(y_true, y_pred)
+
+def soft_dice_loss(y_true, y_pred, axis=(1, 2, 3)):
+    dice_numerator = 2. * K.sum(y_true * y_pred, axis=axis) + 0.00001
+    dice_denominator = K.sum(y_true ** 2, axis=axis) + K.sum(y_pred ** 2, axis=axis) + 0.00001
+    dice_loss = 1 - K.mean((dice_numerator) / (dice_denominator))
+    return dice_loss
+
+
 def find_latest_checkpoint(checkpoints_path, fail_safe=True):
 
     # This is legacy code, there should always be a "checkpoint" file in your directory
@@ -78,6 +97,7 @@ def train(model,
           gen_use_multiprocessing=False,
           ignore_zero_class=False,
           optimizer_name='adam',
+          loss_fun ='categorical_crossentropy',
           do_augment=False,
           augmentation_name="aug_all",
           callbacks=None,
@@ -114,11 +134,11 @@ def train(model,
         if ignore_zero_class:
             loss_k = masked_categorical_crossentropy
         else:
-            loss_k = 'categorical_crossentropy'
+            loss_k = loss_fun #'categorical_crossentropy'
 
         model.compile(loss=loss_k,
                       optimizer=optimizer_name,
-                      metrics=['accuracy'])
+                      metrics=['accuracy',dice_coef])
 
     if checkpoints_path is not None:
         config_file = checkpoints_path + "_config.json"