giou_loss.py 2.63 KB
Newer Older
Hu Ye's avatar
Hu Ye committed
1
import torch
2

3
from ..utils import _log_api_usage_once
4
from ._utils import _loss_inter_union, _upcast_non_float
Hu Ye's avatar
Hu Ye committed
5
6
7
8
9
10
11
12
13


def generalized_box_iou_loss(
    boxes1: torch.Tensor,
    boxes2: torch.Tensor,
    reduction: str = "none",
    eps: float = 1e-7,
) -> torch.Tensor:

Aditya Oke's avatar
Aditya Oke committed
14
    """
Hu Ye's avatar
Hu Ye committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
    Gradient-friendly IoU loss with an additional penalty that is non-zero when the
    boxes do not overlap and scales with the size of their smallest enclosing box.
    This loss is symmetric, so the boxes1 and boxes2 arguments are interchangeable.

    Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
    ``0 <= x1 < x2`` and ``0 <= y1 < y2``, and The two boxes should have the
    same dimensions.

    Args:
        boxes1 (Tensor[N, 4] or Tensor[4]): first set of boxes
        boxes2 (Tensor[N, 4] or Tensor[4]): second set of boxes
        reduction (string, optional): Specifies the reduction to apply to the output:
            ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: No reduction will be
            applied to the output. ``'mean'``: The output will be averaged.
            ``'sum'``: The output will be summed. Default: ``'none'``
Yassine Alouini's avatar
Yassine Alouini committed
30
        eps (float): small number to prevent division by zero. Default: 1e-7
Hu Ye's avatar
Hu Ye committed
31

Aditya Oke's avatar
Aditya Oke committed
32
33
34
    Returns:
        Tensor: Loss tensor with the reduction option applied.

Hu Ye's avatar
Hu Ye committed
35
36
37
38
39
    Reference:
        Hamid Rezatofighi et. al: Generalized Intersection over Union:
        A Metric and A Loss for Bounding Box Regression:
        https://arxiv.org/abs/1902.09630
    """
Aditya Oke's avatar
Aditya Oke committed
40
41
42

    # Original implementation from https://github.com/facebookresearch/fvcore/blob/bfff2ef/fvcore/nn/giou_loss.py

43
44
    if not torch.jit.is_scripting() and not torch.jit.is_tracing():
        _log_api_usage_once(generalized_box_iou_loss)
Hu Ye's avatar
Hu Ye committed
45

Aditya Oke's avatar
Aditya Oke committed
46
47
48
49
50
    boxes1 = _upcast_non_float(boxes1)
    boxes2 = _upcast_non_float(boxes2)
    intsctk, unionk = _loss_inter_union(boxes1, boxes2)
    iouk = intsctk / (unionk + eps)

Hu Ye's avatar
Hu Ye committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    x1, y1, x2, y2 = boxes1.unbind(dim=-1)
    x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1)

    # smallest enclosing box
    xc1 = torch.min(x1, x1g)
    yc1 = torch.min(y1, y1g)
    xc2 = torch.max(x2, x2g)
    yc2 = torch.max(y2, y2g)

    area_c = (xc2 - xc1) * (yc2 - yc1)
    miouk = iouk - ((area_c - unionk) / (area_c + eps))

    loss = 1 - miouk

65
66
67
68
    # Check reduction option and return loss accordingly
    if reduction == "none":
        pass
    elif reduction == "mean":
Hu Ye's avatar
Hu Ye committed
69
70
71
        loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum()
    elif reduction == "sum":
        loss = loss.sum()
72
73
74
75
    else:
        raise ValueError(
            f"Invalid Value for arg 'reduction': '{reduction} \n Supported reduction modes: 'none', 'mean', 'sum'"
        )
Hu Ye's avatar
Hu Ye committed
76
    return loss