From 0df0d6d5cb365238432de2793d92652c82f4a9f6 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 12 Mar 2021 08:26:14 -0800 Subject: [PATCH] fix: transpose batches for tf msssim --- core/leras/ops/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index 398679c..6824488 100644 --- a/core/leras/ops/__init__.py +++ b/core/leras/ops/__init__.py @@ -325,7 +325,10 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value= 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_val=max_value, power_factors=power_factors, + # Transpose images from NCHW to NHWC + img1_t = tf.transpose(img1, [0, 2, 3, 1]) + img2_t = tf.transpose(img2, [0, 2, 3, 1]) + ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors, filter_size=kernel_size, k1=k1, k2=k2) loss = (1.0 - ms_ssim_val) / 2.0