giou_loss.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. import torch
  2. from ..utils import _log_api_usage_once
  3. from ._utils import _loss_inter_union, _upcast_non_float
  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. # Check reduction option and return loss accordingly
  50. if reduction == "none":
  51. pass
  52. elif reduction == "mean":
  53. loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum()
  54. elif reduction == "sum":
  55. loss = loss.sum()
  56. else:
  57. raise ValueError(
  58. f"Invalid Value for arg 'reduction': '{reduction} \n Supported reduction modes: 'none', 'mean', 'sum'"
  59. )
  60. return loss