Intersection over Union Metric/Loss

CHANGELOG.md

@@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added GIoU loss ([#347](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/347))
 
+- Added IoU loss ([#469](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/469))
+
 ### Changed
 
 - Decoupled datamodules from models ([#332](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/332), [#270](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/270))

docs/source/losses.rst

@@ -27,6 +27,14 @@ GIoU Loss
 
 ---------------
 
+IoU Loss
+--------
+
+.. autofunction:: pl_bolts.losses.object_detection.iou_loss
+    :noindex:
+
+---------------
+
 Reinforcement Learning
 ======================
 These are common losses used in RL.

pl_bolts/losses/object_detection.py

@@ -1,10 +1,34 @@
 """
-Generalized Intersection over Union (GIoU) loss (Rezatofighi et. al)
+Loss functions for Object Detection task
 """
 
 import torch
 
-from pl_bolts.metrics.object_detection import giou
+from pl_bolts.metrics.object_detection import giou, iou
+
+
+def iou_loss(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates the intersection over union loss.
+
+    Args:
+        preds: batch of prediction bounding boxes with representation ``[x_min, y_min, x_max, y_max]``
+        target: batch of target bounding boxes with representation ``[x_min, y_min, x_max, y_max]``
+
+    Example:
+
+        >>> import torch
+        >>> from pl_bolts.losses.object_detection import iou_loss
+        >>> preds = torch.tensor([[100, 100, 200, 200]])
+        >>> target = torch.tensor([[150, 150, 250, 250]])
+        >>> iou_loss(preds, target)
+        tensor([[0.8571]])
+
+    Returns:
+        IoU loss
+    """
+    loss = 1 - iou(preds, target)
+    return loss
 
 
 def giou_loss(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:

pl_bolts/metrics/object_detection.py

@@ -1,6 +1,38 @@
 import torch
 
 
+def iou(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates the intersection over union.
+
+    Args:
+        preds: batch of prediction bounding boxes with representation ``[x_min, y_min, x_max, y_max]``
+        target: batch of target bounding boxes with representation ``[x_min, y_min, x_max, y_max]``
+
+    Example:
+
+        >>> import torch
+        >>> from pl_bolts.metrics.object_detection import iou
+        >>> preds = torch.tensor([[100, 100, 200, 200]])
+        >>> target = torch.tensor([[150, 150, 250, 250]])
+        >>> iou(preds, target)
+        tensor([[0.1429]])
+
+    Returns:
+        IoU value
+    """
+    x_min = torch.max(preds[:, None, 0], target[:, 0])
+    y_min = torch.max(preds[:, None, 1], target[:, 1])
+    x_max = torch.min(preds[:, None, 2], target[:, 2])
+    y_max = torch.min(preds[:, None, 3], target[:, 3])
+    intersection = (x_max - x_min).clamp(min=0) * (y_max - y_min).clamp(min=0)
+    pred_area = (preds[:, 2] - preds[:, 0]) * (preds[:, 3] - preds[:, 1])
+    target_area = (target[:, 2] - target[:, 0]) * (target[:, 3] - target[:, 1])
+    union = pred_area[:, None] + target_area - intersection
+    iou = torch.true_divide(intersection, union)
+    return iou
+
+
 def giou(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
     """
     Calculates the generalized intersection over union.

tests/losses/test_object_detection.py

@@ -5,7 +5,22 @@
 import pytest
 import torch
 
-from pl_bolts.losses.object_detection import giou_loss
+from pl_bolts.losses.object_detection import giou_loss, iou_loss
+
+
+@pytest.mark.parametrize("preds, target, expected_loss", [
+    (torch.tensor([[100, 100, 200, 200]]), torch.tensor([[100, 100, 200, 200]]), torch.tensor([0.0]))
+])
+def test_iou_complete_overlap(preds, target, expected_loss):
+    torch.testing.assert_allclose(iou_loss(preds, target), expected_loss)
+
+
+@pytest.mark.parametrize("preds, target, expected_loss", [
+    (torch.tensor([[100, 100, 200, 200]]), torch.tensor([[100, 200, 200, 300]]), torch.tensor([1.0])),
+    (torch.tensor([[100, 100, 200, 200]]), torch.tensor([[200, 200, 300, 300]]), torch.tensor([1.0])),
+])
+def test_iou_no_overlap(preds, target, expected_loss):
+    torch.testing.assert_allclose(iou_loss(preds, target), expected_loss)
 
 
 @pytest.mark.parametrize("preds, target, expected_loss", [

tests/metrics/test_object_detection.py

@@ -5,7 +5,22 @@
 import pytest
 import torch
 
-from pl_bolts.metrics.object_detection import giou
+from pl_bolts.metrics.object_detection import giou, iou
+
+
+@pytest.mark.parametrize("preds, target, expected_iou", [
+    (torch.tensor([[100, 100, 200, 200]]), torch.tensor([[100, 100, 200, 200]]), torch.tensor([1.0]))
+])
+def test_iou_complete_overlap(preds, target, expected_iou):
+    torch.testing.assert_allclose(iou(preds, target), expected_iou)
+
+
+@pytest.mark.parametrize("preds, target, expected_iou", [
+    (torch.tensor([[100, 100, 200, 200]]), torch.tensor([[100, 200, 200, 300]]), torch.tensor([0.0])),
+    (torch.tensor([[100, 100, 200, 200]]), torch.tensor([[200, 200, 300, 300]]), torch.tensor([0.0])),
+])
+def test_iou_no_overlap(preds, target, expected_iou):
+    torch.testing.assert_allclose(iou(preds, target), expected_iou)
 
 
 @pytest.mark.parametrize("preds, target, expected_giou", [