compute ms ssim in test steps

ml4h/metrics.py

@@ -835,15 +835,14 @@ def reset_state(self):
 
 
 class MultiScaleSSIM(keras.metrics.Metric):
-    def __init__(self, max_val=6.0, name="multi_scale_ssim", **kwargs):
+    def __init__(self, name="multi_scale_ssim", **kwargs):
         super(MultiScaleSSIM, self).__init__(name=name, **kwargs)
-        self.max_val = max_val
         self.total_ssim = self.add_weight(name="total_ssim", initializer="zeros")
         self.count = self.add_weight(name="count", initializer="zeros")
 
-    def update_state(self, y_true, y_pred, sample_weight=None):
+    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=self.max_val)
+        ssim = tf.image.ssim_multiscale(y_true, y_pred, max_val=max_val)
         if sample_weight is not None:
             ssim = tf.multiply(ssim, sample_weight)
 
@@ -855,7 +854,7 @@ def result(self):
         # Return the mean MS-SSIM over all batches
         return tf.divide(self.total_ssim, self.count)
 
-    def reset_states(self):
+    def reset_state(self):
         # Reset the metric state variables
         self.total_ssim.assign(0.0)
         self.count.assign(0.0)

ml4h/models/diffusion_blocks.py

@@ -886,7 +886,9 @@ def test_step(self, batch):
                 num_images=self.batch_size, diffusion_steps=20
             )
             self.kid.update_state(images, generated_images)
-            self.ms_ssim.update_state(images, generated_images)
+            max_pixel_value = tf.reduce_max(tf.abs(images))
+            max_val = 2 * max_pixel_value.numpy()  # Double the max absolute value
+            self.ms_ssim.update_state(images, generated_images, max_val)
 
         return {m.name: m.result() for m in self.metrics}