feat: ms_ssim implementation

core/leras/ops/__init__.py

@@ -307,6 +307,35 @@ def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03
 
 nn.dssim = dssim
 
+
+def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=1.0,
+            power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)):
+    img_dtype = img1.dtype
+    if img_dtype != img2.dtype:
+        raise ValueError("img1.dtype != img2.dtype")
+
+    if img_dtype != tf.float32:
+        img1 = tf.cast(img1, tf.float32)
+        img2 = tf.cast(img2, tf.float32)
+
+    # restrict mssim factors to those greater/equal to kernel size
+    power_factors = [power_factors[i] for i in range(len(power_factors)) if resolution//(2**i) >= kernel_size]
+
+    # normalize power factors if reduced because of size
+    if sum(power_factors) < 1.0:
+        power_factors = [x/sum(power_factors) for x in power_factors]
+
+    ms_ssim_val = tf.image.ssim_multiscale(img1, img2, max_value=max_value, power_factors=power_factors,
+                                           filter_size=kernel_size, k1=k1, k2=k2)
+    loss = (1.0 - ms_ssim_val) / 2.0
+
+    if img_dtype != tf.float32:
+        loss = tf.cast(loss, img_dtype)
+    return loss
+
+nn.dssim = ms_ssim
+
+
 def space_to_depth(x, size):
     if nn.data_format == "NHWC":
         # match NCHW version in order to switch data_format without problems