add ssim evaluation metric

video_prediction_tools/main_scripts/main_visualize_postprocess.py

@@ -14,16 +14,12 @@ import pandas as pd
 import tensorflow as tf
 import warnings
 import pickle
-from random import seed
 import datetime
 import json
-from netCDF4 import Dataset, date2num
 import matplotlib
-
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 from mpl_toolkits.basemap import Basemap
-import matplotlib.gridspec as gridspec
 from normalization import Norm_data
 from metadata import MetaData as MetaData
 from main_scripts.main_train_models import *

video_prediction_tools/model_modules/video_prediction/metrics.py

@@ -2,6 +2,8 @@ import tensorflow as tf
 #import lpips_tf
 import numpy as np
 import math
+from skimage.metrics import structural_similarity as ssim
+
 def mse(a, b):
     return tf.reduce_mean(tf.squared_difference(a, b), [-3, -2, -1])
 
@@ -24,6 +26,7 @@ def psnr_imgs(img1, img2):
 def mse_imgs(image1,image2):
     mse = ((image1 - image2)**2).mean(axis=None)
     return mse
+
 # def lpips(input0, input1):
 #     if input0.shape[-1].value == 1:
 #         input0 = tf.tile(input0, [1] * (input0.shape.ndims - 1) + [3])
@@ -35,7 +38,12 @@ def mse_imgs(image1,image2):
 
 def ssim_images(image1, image2):
     """
-
+    Reference for calculating ssim
     Numpy impelmeentation for ssim https://cvnote.ddlee.cc/2019/09/12/psnr-ssim-python
+    https://scikit-image.org/docs/dev/auto_examples/transform/plot_ssim.html
+    :param image1 the reference images
+    :param image2 the predicte images
     """
-    pass    
+    ssim_pred = ssim(image1, image2,
+                      data_range = image2.max() - image2.min())
+    return ssim_pred