compute ms ssim in test steps

broadinstitute · Jan 21, 2025 · a6f26f2 · a6f26f2
1 parent a65f88d
commit a6f26f2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/ml4h/metrics.py b/ml4h/metrics.py
@@ -842,7 +842,7 @@ def __init__(self, name="multi_scale_ssim", **kwargs):
 
     def update_state(self, y_true, y_pred, max_val, sample_weight=None):
         # Calculate MS-SSIM for the batch
-        ssim = tf.image.ssim_multiscale(y_true, y_pred, max_val=max_val)
+        ssim = tf.image.ssim_multiscale(y_true, y_pred, max_val=max_val, power_factors=[0.01, 0.2, 0.5, 0.29])
         if sample_weight is not None:
             ssim = tf.multiply(ssim, sample_weight)
 

diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
@@ -888,8 +888,7 @@ def test_step(self, batch):
             self.kid.update_state(images, generated_images)
             max_pixel_value = tf.reduce_max(tf.abs(images))
             max_val = 2 * max_pixel_value  # Double the max absolute value
-            self.ms_ssim.update_state(tf.tile(images, [1, 1, 1, 3]),
-                                      tf.tile(generated_images, [1, 1, 1, 3]), max_val)
+            self.ms_ssim.update_state(images, generated_images, max_val)
 
         return {m.name: m.result() for m in self.metrics}