This commit is contained in:
jh 2021-03-17 11:15:33 -07:00
commit 33ea98012c

View file

@ -1,6 +1,7 @@
from core.leras import nn from core.leras import nn
tf = nn.tf tf = nn.tf
class MsSsim(nn.LayerBase): class MsSsim(nn.LayerBase):
default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
@ -20,7 +21,9 @@ class MsSsim(nn.LayerBase):
y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1])
y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1])
loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size) ms_ssim_val = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size)
return (1.0 - loss) / 2.0 # ssim_multiscale returns values in range [0, 1] (where 1 is completely identical)
# subtract from 1 to get loss
return 1 - ms_ssim_val
nn.MsSsim = MsSsim nn.MsSsim = MsSsim