Speeding up mean_iou metric computation (#569)

* Change Features from Sequence to Image * Fix example
huggingface · Apr 18, 2024 · edb83af · edb83af
1 parent 04c7588
commit edb83af
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/metrics/mean_iou/README.md b/metrics/mean_iou/README.md
@@ -27,13 +27,13 @@ For binary (two classes) or multi-class segmentation, the *mean IoU* of the imag
 
 ## How to Use
 
-The Mean IoU metric takes two numeric arrays as input corresponding to the predicted and ground truth segmentations:
+The Mean IoU metric takes two lists of numeric 2D arrays as input corresponding to the predicted and ground truth segmentations:
 ```python
 >>> import numpy as np
 >>> mean_iou = evaluate.load("mean_iou")
 >>> predicted = np.array([[2, 2, 3], [8, 2, 4], [3, 255, 2]])
 >>> ground_truth = np.array([[1, 2, 2], [8, 2, 1], [3, 255, 1]])
->>> results = mean_iou.compute(predictions=predicted, references=ground_truth, num_labels=10, ignore_index=255)
+>>> results = mean_iou.compute(predictions=[predicted], references=[ground_truth], num_labels=10, ignore_index=255)
 ```
 
 ### Inputs

diff --git a/metrics/mean_iou/mean_iou.py b/metrics/mean_iou/mean_iou.py
@@ -281,10 +281,9 @@ def _info(self):
             citation=_CITATION,
             inputs_description=_KWARGS_DESCRIPTION,
             features=datasets.Features(
-                # 1st Seq - height dim, 2nd - width dim
                 {
-                    "predictions": datasets.Sequence(datasets.Sequence(datasets.Value("uint16"))),
-                    "references": datasets.Sequence(datasets.Sequence(datasets.Value("uint16"))),
+                    "predictions": datasets.Image(),
+                    "references": datasets.Image(),
                 }
             ),
             reference_urls=[