[metrics] IoU Metric (#2062)

* add iou function

* update stat scores

* add iou class

* add iou tests

* chlog

* Apply suggestions from code review

* tests

* docs

* Apply suggestions from code review

* docs

Co-authored-by: Jirka <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>

Author: j-dsouza

CHANGELOG.md

@@ -13,7 +13,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added metrics 
   * Base classes ([#1326](https://github.com/PyTorchLightning/pytorch-lightning/pull/1326), [#1877](https://github.com/PyTorchLightning/pytorch-lightning/pull/1877))
   * Sklearn metrics classes ([#1327](https://github.com/PyTorchLightning/pytorch-lightning/pull/1327))
-  * Native torch metrics ([#1488](https://github.com/PyTorchLightning/pytorch-lightning/pull/1488))
+  * Native torch metrics ([#1488](https://github.com/PyTorchLightning/pytorch-lightning/pull/1488), [#2062](https://github.com/PyTorchLightning/pytorch-lightning/pull/2062))
   * docs for all Metrics ([#2184](https://github.com/PyTorchLightning/pytorch-lightning/pull/2184), [#2209](https://github.com/PyTorchLightning/pytorch-lightning/pull/2209))
   * Regression metrics ([#2221](https://github.com/PyTorchLightning/pytorch-lightning/pull/2221))
 - Added type hints in `Trainer.fit()` and `Trainer.test()` to reflect that also a list of dataloaders can be passed in ([#1723](https://github.com/PyTorchLightning/pytorch-lightning/pull/1723))

docs/source/metrics.rst

@@ -31,7 +31,7 @@ Example::
     to a few metrics. Please feel free to create an issue/PR if you have a proposed 
     metric or have found a bug.
 
---------------
+---
 
 Implement a metric
 ------------------
@@ -75,7 +75,7 @@ Here's an example showing how to implement a NumpyMetric
 .. autoclass:: pytorch_lightning.metrics.metric.NumpyMetric
     :noindex:
 
---------------
+---
 
 Class Metrics
 -------------
@@ -207,6 +207,12 @@ MulticlassPrecisionRecall
 .. autoclass:: pytorch_lightning.metrics.classification.MulticlassPrecisionRecall
     :noindex:
 
+IoU
+^^^
+
+.. autoclass:: pytorch_lightning.metrics.classification.IoU
+    :noindex:
+
 RMSE
 ^^^^
 
@@ -219,7 +225,7 @@ RMSLE
 .. autoclass:: pytorch_lightning.metrics.regression.RMSE
     :noindex:
 
---------------
+---
 
 Functional Metrics
 ------------------
@@ -346,16 +352,23 @@ stat_scores (F)
 .. autofunction:: pytorch_lightning.metrics.functional.stat_scores
     :noindex:
 
+iou (F)
+^^^^^^^
+
+.. autofunction:: pytorch_lightning.metrics.functional.iou
+    :noindex:
+
 stat_scores_multiple_classes (F)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. autofunction:: pytorch_lightning.metrics.functional.stat_scores_multiple_classes
     :noindex:
 
-----------------
+---
 
 Metric pre-processing
 ---------------------
+
 Metric
 
 to_categorical (F)
@@ -370,7 +383,7 @@ to_onehot (F)
 .. autofunction:: pytorch_lightning.metrics.functional.to_onehot
     :noindex:
 
-----------------
+---
 
 Sklearn interface
 -----------------

pytorch_lightning/metrics/__init__.py

@@ -20,6 +20,7 @@
     MulticlassROC,
     Precision,
     PrecisionRecall,
+    IoU,
 )
 from pytorch_lightning.metrics.sklearns import (
     AUC,
@@ -43,6 +44,7 @@
     'PrecisionRecallCurve',
     'ROC',
     'Recall',
+    'IoU',
 ]
 __regression_metrics = [
     'MSE',

pytorch_lightning/metrics/classification.py

@@ -15,7 +15,8 @@
     roc,
     multiclass_roc,
     multiclass_precision_recall_curve,
-    dice_score
+    dice_score,
+    iou,
 )
 from pytorch_lightning.metrics.metric import TensorMetric, TensorCollectionMetric
 
@@ -770,3 +771,48 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                           nan_score=self.nan_score,
                           no_fg_score=self.no_fg_score,
                           reduction=self.reduction)
+
+
+class IoU(TensorMetric):
+    """
+    Computes the intersection over union.
+
+    Example:
+
+        >>> pred = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 0],
+        ...                      [0, 0, 1, 1, 1, 0, 0, 0],
+        ...                      [0, 0, 0, 0, 0, 0, 0, 0]])
+        >>> target = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 0],
+        ...                        [0, 0, 0, 1, 1, 1, 0, 0],
+        ...                        [0, 0, 0, 0, 0, 0, 0, 0]])
+        >>> metric = IoU()
+        >>> metric(pred, target)
+        tensor(0.7045)
+
+    """
+    def __init__(self,
+                 remove_bg: bool = False,
+                 reduction: str = 'elementwise_mean'):
+        """
+        Args:
+            remove_bg: Flag to state whether a background class has been included
+                within input parameters. If true, will remove background class. If
+                false, return IoU over all classes.
+                Assumes that background is '0' class in input tensor
+            reduction: a method for reducing IoU over labels (default: takes the mean)
+                Available reduction methods:
+
+                - elementwise_mean: takes the mean
+                - none: pass array
+                - sum: add elements
+        """
+        super().__init__(name='iou')
+        self.remove_bg = remove_bg
+        self.reduction = reduction
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor,
+                sample_weight: Optional[torch.Tensor] = None):
+        """
+        Actual metric calculation.
+        """
+        return iou(y_pred, y_true, remove_bg=self.remove_bg, reduction=self.reduction)

pytorch_lightning/metrics/functional/__init__.py

@@ -17,5 +17,6 @@
     stat_scores,
     stat_scores_multiple_classes,
     to_categorical,
-    to_onehot
+    to_onehot,
+    iou,
 )

pytorch_lightning/metrics/functional/classification.py

@@ -901,3 +901,49 @@ def dice_score(
 
         scores[i - bg] += score_cls
     return reduce(scores, reduction=reduction)
+
+
+def iou(pred: torch.Tensor, target: torch.Tensor,
+        num_classes: Optional[int] = None, remove_bg: bool = False,
+        reduction: str = 'elementwise_mean'):
+    """
+    Intersection over union, or Jaccard index calculation.
+
+    Args:
+        pred: Tensor containing predictions
+
+        target: Tensor containing targets
+
+        num_classes: Optionally specify the number of classes
+
+        remove_bg: Flag to state whether a background class has been included
+            within input parameters. If true, will remove background class. If
+            false, return IoU over all classes.
+            Assumes that background is '0' class in input tensor
+
+        reduction: a method for reducing IoU over labels (default: takes the mean)
+            Available reduction methods:
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum: add elements
+
+    Returns:
+        IoU score : Tensor containing single value if reduction is
+        'elementwise_mean', or number of classes if reduction is 'none'
+
+    Example:
+
+        >>> target = torch.randint(0, 1, (10, 25, 25))
+        >>> pred = torch.tensor(target)
+        >>> pred[2:5, 7:13, 9:15] = 1 - pred[2:5, 7:13, 9:15]
+        >>> iou(pred, target)
+        tensor(0.4914)
+
+    """
+    tps, fps, tns, fns, sups = stat_scores_multiple_classes(pred, target, num_classes)
+    if remove_bg:
+        tps = tps[1:]
+        fps = fps[1:]
+        fns = fns[1:]
+    iou = tps / (fps + fns + tps)
+    return reduce(iou, reduction=reduction)

tests/metrics/functional/test_classification.py

@@ -31,6 +31,7 @@
     precision_recall_curve,
     roc,
     auc,
+    iou,
 )
 
 
@@ -366,5 +367,22 @@ def test_dice_score(pred, target, expected):
     assert score == expected
 
 
+@pytest.mark.parametrize(['half_ones', 'reduction', 'remove_bg', 'expected'], [
+    pytest.param(False, 'none', False, torch.Tensor([1, 1, 1])),
+    pytest.param(False, 'elementwise_mean', False, torch.Tensor([1])),
+    pytest.param(False, 'none', True, torch.Tensor([1, 1])),
+    pytest.param(True, 'none', False, torch.Tensor([0.5, 0.5, 0.5])),
+    pytest.param(True, 'elementwise_mean', False, torch.Tensor([0.5])),
+    pytest.param(True, 'none', True, torch.Tensor([0.5, 0.5])),
+])
+def test_iou(half_ones, reduction, remove_bg, expected):
+    pred = (torch.arange(120) % 3).view(-1, 1)
+    target = (torch.arange(120) % 3).view(-1, 1)
+    if half_ones:
+        pred[:60] = 1
+    iou_val = iou(pred, target, remove_bg=remove_bg, reduction=reduction)
+    assert torch.allclose(iou_val, expected, atol=1e-9)
+
+
 # example data taken from
 # https://github.com/scikit-learn/scikit-learn/blob/master/sklearn/metrics/tests/test_ranking.py

tests/metrics/test_classification.py

@@ -19,6 +19,7 @@
     MulticlassROC,
     MulticlassPrecisionRecall,
     DiceCoefficient,
+    IoU,
 )
 
 
@@ -205,3 +206,14 @@ def test_dice_coefficient(include_background):
     dice = dice_coeff(torch.randint(0, 1, (10, 25, 25)),
                       torch.randint(0, 1, (10, 25, 25)))
     assert isinstance(dice, torch.Tensor)
+
+
+@pytest.mark.parametrize('remove_bg', [True, False])
+def test_iou(remove_bg):
+    iou = IoU(remove_bg=remove_bg)
+    assert iou.name == 'iou'
+
+    score = iou(torch.randint(0, 1, (10, 25, 25)),
+                torch.randint(0, 1, (10, 25, 25)))
+
+    assert isinstance(score, torch.Tensor)