giou_loss.py 2.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. import torch
  2. from ..utils import _log_api_usage_once
  3. from ._utils import _upcast_non_float, _loss_inter_union
  4. def generalized_box_iou_loss(
  5. boxes1: torch.Tensor,
  6. boxes2: torch.Tensor,
  7. reduction: str = "none",
  8. eps: float = 1e-7,
  9. ) -> torch.Tensor:
  10. """
  11. Gradient-friendly IoU loss with an additional penalty that is non-zero when the
  12. boxes do not overlap and scales with the size of their smallest enclosing box.
  13. This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable.
  14. Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
  15. ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the
  16. same dimensions.
  17. Args:
  18. boxes1 (Tensor[N, 4] or Tensor[4]): first set of boxes
  19. boxes2 (Tensor[N, 4] or Tensor[4]): second set of boxes
  20. reduction (string, optional): Specifies the reduction to apply to the output:
  21. ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be
  22. applied to the output. ``'mean'``: The output will be averaged.
  23. ``'sum'``: The output will be summed. Default: ``'none'``
  24. eps (float): small number to prevent division by zero. Default: 1e-7
  25. Returns:
  26. Tensor: Loss tensor with the reduction option applied.
  27. Reference:
  28. Hamid Rezatofighi et. al: Generalized Intersection over Union:
  29. A Metric and A Loss for Bounding Box Regression:
  30. https://arxiv.org/abs/1902.09630
  31. """
  32. # Original implementation from https://github.com/facebookresearch/fvcore/blob/bfff2ef/fvcore/nn/giou_loss.py
  33. if not torch.jit.is_scripting() and not torch.jit.is_tracing():
  34. _log_api_usage_once(generalized_box_iou_loss)
  35. boxes1 = _upcast_non_float(boxes1)
  36. boxes2 = _upcast_non_float(boxes2)
  37. intsctk, unionk = _loss_inter_union(boxes1, boxes2)
  38. iouk = intsctk / (unionk + eps)
  39. x1, y1, x2, y2 = boxes1.unbind(dim=-1)
  40. x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1)
  41. # smallest enclosing box
  42. xc1 = torch.min(x1, x1g)
  43. yc1 = torch.min(y1, y1g)
  44. xc2 = torch.max(x2, x2g)
  45. yc2 = torch.max(y2, y2g)
  46. area_c = (xc2 - xc1) * (yc2 - yc1)
  47. miouk = iouk - ((area_c - unionk) / (area_c + eps))
  48. loss = 1 - miouk
  49. if reduction == "mean":
  50. loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum()
  51. elif reduction == "sum":
  52. loss = loss.sum()
  53. return loss