Use SSIM loss function with Keras

Question

I need to use the SSIM from Sewar as a loss function in order to compare images for my model.

I am getting errors when I try to compile my model. I import the function and compile the model like this:

from sewar.full_ref import ssim
...
model.compile('ssim', optimizer=my_optimizer, metrics=[ssim])

and I get this:

File "/media/merry/merry32/train.py", line 19, in train
model.compile(loss='ssim', optimizer=opt, metrics=[ssim])
File "/home/merry/anaconda3/envs/merry_env/lib/python3.7/site-packages/keras/engine/training.py", line 451, in compile
handle_metrics(output_metrics)
File "/home/merry/anaconda3/envs/merry_env/lib/python3.7/site-packages/keras/engine/training.py", line 420, in handle_metrics
mask=masks[i])
File "/home/merry/anaconda3/envs/merry_env/lib/python3.7/site-packages/keras/engine/training_utils.py", line 404, in weighted
score_array = fn(y_true, y_pred)
File "/home/merry/anaconda3/envs/merry_env/lib/python3.7/site-packages/sewar/full_ref.py", line 143, in ssim
MAX = np.iinfo(GT.dtype).max
File "/home/merry/anaconda3/envs/merry_env/lib/python3.7/site-packages/numpy/core/getlimits.py", line 506, in __init__
raise ValueError("Invalid integer data type %r." % (self.kind,))
ValueError: Invalid integer data type 'O'.

I could also write something like this:

model.compile(ssim(), optimizer=my_optimizer, metrics=[ssim()])

But then I get this error (obviously):

TypeError: ssim() missing 2 required positional arguments: 'GT' and 'P'

I just wanted to do the same I was doing with mean_sqeared_error but with SSIM, like this (which works perfectly with no need of passing parameters to it):

model.compile('mean_squared_error', optimizer=my_optimizer, metrics=['mse'])

Any idea on how should I use this function to compile?

Answer 1

Keras has an implementation of SSIM. You can use it like this:

def SSIMLoss(y_true, y_pred):
  return 1 - tf.reduce_mean(tf.image.ssim(y_true, y_pred, 1.0))

self.model.compile(optimizer=sgd, loss=SSIMLoss)

Answer 2

• You can use tf.image.ssim to compute SSIM index between two images.
• Since training happens on batch of images we will use the mean of SSIM values of all the images in the batch as the loss value
• Our model will return an image (of some size based on the CNN layers used which is again based on input and expected output image dimensions).

Sample working code

from keras.models import Sequential
from keras.layers import Dense, Conv2D, Flatten
import numpy as np
import tensorflow as tf

# Loss functtion
def ssim_loss(y_true, y_pred):
  return tf.reduce_mean(tf.image.ssim(y_true, y_pred, 2.0))

# Model: Input Image size: 32X32X1 output Image size: 28X28X1 
# check model.summary
model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3),
                 activation='relu',
                 input_shape=(32,32,1)))
model.add(Conv2D(1, kernel_size=(3, 3),
                 activation='relu'))

model.compile(optimizer='adam', loss=ssim_loss, metrics=[ssim_loss, 'accuracy'])

# Train
model.fit(np.random.randn(10,32,32,1), np.random.randn(10,28,28,1))