add Dice to metrax (#64) (#91)

* add Dice to metrax (#64)

* fix: added sample_weights parameter for DICE

* removed sample_weights from DICE

* moved DICE form classification_metrics.py to image_metrics.py

Author: Mrigankkh

src/metrax/__init__.py

@@ -26,6 +26,7 @@
 AveragePrecisionAtK = ranking_metrics.AveragePrecisionAtK
 BLEU = nlp_metrics.BLEU
 DCGAtK = ranking_metrics.DCGAtK
+Dice = image_metrics.Dice
 IoU = image_metrics.IoU
 MAE = regression_metrics.MAE
 MRR = ranking_metrics.MRR
@@ -53,6 +54,7 @@
     "AveragePrecisionAtK",
     "BLEU",
     "DCGAtK",
+    "Dice",
     "IoU",
     "MAE",
     "MRR",

src/metrax/classification_metrics.py

@@ -576,4 +576,4 @@ def compute(self) -> jax.Array:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.
-    return jnp.trapezoid(tp_rate, fp_rate) * -1
+    return jnp.trapezoid(tp_rate, fp_rate) * -1

src/metrax/classification_metrics_test.py

@@ -47,7 +47,6 @@
     (BATCHES, 1),
 ).astype(np.float32)
 
-
 class ClassificationMetricsTest(parameterized.TestCase):
 
   def test_precision_empty(self):
@@ -78,7 +77,7 @@ def test_aucroc_empty(self):
     self.assertEqual(m.false_positives, jnp.array(0, jnp.float32))
     self.assertEqual(m.false_negatives, jnp.array(0, jnp.float32))
     self.assertEqual(m.num_thresholds, 0)
-
+    
   @parameterized.named_parameters(
       ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS),
       ('basic_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS),

src/metrax/image_metrics.py

@@ -14,6 +14,7 @@
 
 """A collection of different metrics for image models."""
 
+from clu import metrics as clu_metrics
 import flax
 import jax
 from jax import lax
@@ -584,3 +585,82 @@ def from_model_output(
       """
       batch_psnr = cls._calculate_psnr(predictions, targets, max_val=max_val)
       return super().from_model_output(values=batch_psnr)  
+  
+
+@flax.struct.dataclass
+class Dice(clu_metrics.Metric):
+  r"""Computes the Dice coefficient between `y_true` and `y_pred`.
+
+  Dice is a similarity measure used to measure overlap between two samples.
+  A Dice score of 1 indicates perfect overlap; 0 indicates no overlap.
+
+  The formula is:
+
+  .. math::
+
+      \text{Dice} = \frac{2 \cdot \sum (y_{\text{true}} \cdot y_{\text{pred}})}
+                          {\sum y_{\text{true}} + \sum y_{\text{pred}} + \epsilon}
+
+  Attributes:
+      intersection: Sum of element-wise product between `y_true` and `y_pred`.
+      sum_true: Sum of y_true across all examples.
+      sum_pred: Sum of y_pred across all examples.
+      
+  """
+ 
+  intersection: jax.Array
+  sum_pred: jax.Array
+  sum_true: jax.Array
+  
+  @classmethod
+  def empty(cls) -> 'Dice':
+      return cls(
+          intersection=jnp.array(0.0, jnp.float32),
+          sum_pred=jnp.array(0.0, jnp.float32),
+          sum_true=jnp.array(0.0, jnp.float32),
+      )
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: jax.Array,
+      labels: jax.Array,
+  ) -> 'Dice':
+      """Updates the metric.
+
+      Args:
+          predictions: A floating point vector whose values are in the range [0,
+            1]. The shape should be (batch_size,).
+          labels: True value. The value is expected to be 0 or 1. The shape should
+            be (batch_size,).
+
+      Returns:
+          Updated Dice metric. 
+          
+      """
+      predictions = jnp.asarray(predictions, jnp.float32)
+      labels = jnp.asarray(labels, jnp.float32)
+
+      intersection = jnp.sum(predictions * labels)
+      sum_pred = jnp.sum(predictions)
+      sum_true = jnp.sum(labels)
+
+      return cls(
+          intersection=intersection,
+          sum_pred=sum_pred,
+          sum_true=sum_true,
+      )
+
+  def merge(self, other: 'Dice') -> 'Dice':
+      return type(self)(
+          intersection=self.intersection + other.intersection,
+          sum_pred=self.sum_pred + other.sum_pred,
+          sum_true=self.sum_true + other.sum_true,
+      )
+
+  def compute(self) -> jax.Array:
+      """Returns the final Dice coefficient."""
+      epsilon = 1e-7
+      return (2.0 * self.intersection) / (
+          self.sum_pred + self.sum_true + epsilon
+      )

src/metrax/image_metrics_test.py

@@ -190,9 +190,38 @@
 NUM_CLASSES_IOU_6 = NUM_CLASSES_IOU_2
 TARGET_CLASS_IDS_IOU_6 = np.array(range(NUM_CLASSES_IOU_6))
 
+# Test data for Dice
+BATCHES = 4
+BATCH_SIZE = 8
+OUTPUT_LABELS = np.random.randint(
+    0,
+    2,
+    size=(BATCHES, BATCH_SIZE),
+).astype(np.float32)
+OUTPUT_PREDS = np.random.uniform(size=(BATCHES, BATCH_SIZE))
+OUTPUT_PREDS_F16 = OUTPUT_PREDS.astype(jnp.float16)
+OUTPUT_PREDS_F32 = OUTPUT_PREDS.astype(jnp.float32)
+OUTPUT_PREDS_BF16 = OUTPUT_PREDS.astype(jnp.bfloat16)
+OUTPUT_LABELS_BS1 = np.random.randint(
+    0,
+    2,
+    size=(BATCHES, 1),
+).astype(np.float32)
+OUTPUT_PREDS_BS1 = np.random.uniform(size=(BATCHES, 1)).astype(np.float32)
 
-class ImageMetricsTest(parameterized.TestCase):
+DICE_ALL_ONES = (jnp.array([1, 1, 1, 1]), jnp.array([1, 1, 1, 1]))
+DICE_ALL_ZEROS = (jnp.array([0, 0, 0, 0]), jnp.array([0, 0, 0, 0]))
+DICE_NO_OVERLAP = (jnp.array([1, 1, 0, 0]), jnp.array([0, 0, 1, 1]))
 
+class ImageMetricsTest(parameterized.TestCase):
+    
+  def test_dice_empty(self):
+    """Tests the `empty` method of the `Dice` class."""
+    m = metrax.Dice.empty()
+    self.assertEqual(m.intersection, jnp.array(0, jnp.float32))
+    self.assertEqual(m.sum_true, jnp.array(0, jnp.float32))
+    self.assertEqual(m.sum_pred, jnp.array(0, jnp.float32))
+    
   @parameterized.named_parameters(
       (
           'ssim_basic_norm_single_channel',
@@ -500,5 +529,36 @@ def test_psnr_against_tensorflow(
                 atol=1e-4,
                 err_msg="PSNR mismatch",
     )
+            
+  @parameterized.named_parameters(
+      ('basic_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32),
+      ('low_threshold', OUTPUT_LABELS, OUTPUT_PREDS_F32),
+      ('high_threshold', OUTPUT_LABELS, OUTPUT_PREDS_F32),
+      ('batch_size_one', OUTPUT_LABELS_BS1, OUTPUT_PREDS_BS1),
+      ('all_ones', *DICE_ALL_ONES),  
+      ('all_zeros', *DICE_ALL_ZEROS),
+      ('no_overlap', *DICE_NO_OVERLAP),
+    )             
+  def test_dice(self, y_true, y_pred):
+    """Test that Dice metric computes expected values."""
+    y_true = jnp.asarray(y_true, jnp.float32)
+    y_pred = jnp.asarray(y_pred, jnp.float32)
+
+    # Manually compute expected Dice
+    eps = 1e-7
+    intersection = jnp.sum(y_true * y_pred)
+    sum_pred     = jnp.sum(y_pred)
+    sum_true     = jnp.sum(y_true)
+    expected     = (2.0 * intersection) / (sum_pred + sum_true + eps)
+
+    # Compute using the metric class
+    metric = metrax.Dice.from_model_output(
+        predictions=y_pred,
+        labels=y_true
+    )
+    result = metric.compute()
+    
+    np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-5)
+    
 if __name__ == '__main__':
   absltest.main()

src/metrax/metrax_test.py

@@ -98,6 +98,11 @@ class MetraxTest(parameterized.TestCase):
               'ks': KS,
           },
       ),
+      (
+          'dice',
+          metrax.Dice,
+          {'predictions': OUTPUT_LABELS, 'labels': OUTPUT_PREDS},
+      ),
       (
           'iou',
           metrax.IoU,

src/metrax/nnx/__init__.py

@@ -21,6 +21,7 @@
 AveragePrecisionAtK = nnx_metrics.AveragePrecisionAtK
 BLEU = nnx_metrics.BLEU
 DCGAtK = nnx_metrics.DCGAtK
+Dice = nnx_metrics.Dice
 IoU = nnx_metrics.IoU
 MAE = nnx_metrics.MAE
 MRR = nnx_metrics.MRR
@@ -47,6 +48,7 @@
     "AveragePrecisionAtK",
     "BLEU",
     "DCGAtK",
+    "Dice",
     "IoU",
     "MRR",
     "MAE"

src/metrax/nnx/nnx_metrics.py

@@ -66,6 +66,13 @@ def __init__(self):
     super().__init__(metrax.DCGAtK)
 
 
+class Dice(NnxWrapper):
+  """An NNX class for the Metrax metric Dice."""
+
+  def __init__(self):
+    super().__init__(metrax.Dice)
+
+
 class IoU(NnxWrapper):
   """An NNX class for the Metrax metric IoU."""