|
| 1 | +import math |
| 2 | +import torch |
| 3 | + |
| 4 | + |
| 5 | +def diou_loss( |
| 6 | + boxes1: torch.Tensor, |
| 7 | + boxes2: torch.Tensor, |
| 8 | + reduction: str = "none", |
| 9 | + eps: float = 1e-7, |
| 10 | +) -> torch.Tensor: |
| 11 | + """ |
| 12 | + Distance Intersection over Union Loss (Zhaohui Zheng et. al) |
| 13 | + https://arxiv.org/abs/1911.08287 |
| 14 | + Args: |
| 15 | + boxes1, boxes2 (Tensor): box locations in XYXY format, shape (N, 4) or (4,). |
| 16 | + reduction: 'none' | 'mean' | 'sum' |
| 17 | + 'none': No reduction will be applied to the output. |
| 18 | + 'mean': The output will be averaged. |
| 19 | + 'sum': The output will be summed. |
| 20 | + eps (float): small number to prevent division by zero |
| 21 | + """ |
| 22 | + |
| 23 | + x1, y1, x2, y2 = boxes1.unbind(dim=-1) |
| 24 | + x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) |
| 25 | + |
| 26 | + # TODO: use torch._assert_async() when pytorch 1.8 support is dropped |
| 27 | + assert (x2 >= x1).all(), "bad box: x1 larger than x2" |
| 28 | + assert (y2 >= y1).all(), "bad box: y1 larger than y2" |
| 29 | + |
| 30 | + # Intersection keypoints |
| 31 | + xkis1 = torch.max(x1, x1g) |
| 32 | + ykis1 = torch.max(y1, y1g) |
| 33 | + xkis2 = torch.min(x2, x2g) |
| 34 | + ykis2 = torch.min(y2, y2g) |
| 35 | + |
| 36 | + intsct = torch.zeros_like(x1) |
| 37 | + mask = (ykis2 > ykis1) & (xkis2 > xkis1) |
| 38 | + intsct[mask] = (xkis2[mask] - xkis1[mask]) * (ykis2[mask] - ykis1[mask]) |
| 39 | + union = (x2 - x1) * (y2 - y1) + (x2g - x1g) * (y2g - y1g) - intsct + eps |
| 40 | + iou = intsct / union |
| 41 | + |
| 42 | + # smallest enclosing box |
| 43 | + xc1 = torch.min(x1, x1g) |
| 44 | + yc1 = torch.min(y1, y1g) |
| 45 | + xc2 = torch.max(x2, x2g) |
| 46 | + yc2 = torch.max(y2, y2g) |
| 47 | + diag_len = ((xc2 - xc1) ** 2) + ((yc2 - yc1) ** 2) + eps |
| 48 | + |
| 49 | + # centers of boxes |
| 50 | + x_p = (x2 + x1) / 2 |
| 51 | + y_p = (y2 + y1) / 2 |
| 52 | + x_g = (x1g + x2g) / 2 |
| 53 | + y_g = (y1g + y2g) / 2 |
| 54 | + distance = ((x_p - x_g) ** 2) + ((y_p - y_g) ** 2) |
| 55 | + |
| 56 | + # Eqn. (7) |
| 57 | + loss = 1 - iou + (distance / diag_len) |
| 58 | + if reduction == "mean": |
| 59 | + loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() |
| 60 | + elif reduction == "sum": |
| 61 | + loss = loss.sum() |
| 62 | + |
| 63 | + return loss |
| 64 | + |
| 65 | + |
| 66 | +def ciou_loss( |
| 67 | + boxes1: torch.Tensor, |
| 68 | + boxes2: torch.Tensor, |
| 69 | + reduction: str = "none", |
| 70 | + eps: float = 1e-7, |
| 71 | +) -> torch.Tensor: |
| 72 | + """ |
| 73 | + Complete Intersection over Union Loss (Zhaohui Zheng et. al) |
| 74 | + https://arxiv.org/abs/1911.08287 |
| 75 | + Args: |
| 76 | + boxes1, boxes2 (Tensor): box locations in XYXY format, shape (N, 4) or (4,). |
| 77 | + reduction: 'none' | 'mean' | 'sum' |
| 78 | + 'none': No reduction will be applied to the output. |
| 79 | + 'mean': The output will be averaged. |
| 80 | + 'sum': The output will be summed. |
| 81 | + eps (float): small number to prevent division by zero |
| 82 | + """ |
| 83 | + |
| 84 | + x1, y1, x2, y2 = boxes1.unbind(dim=-1) |
| 85 | + x1g, y1g, x2g, y2g = boxes2.unbind(dim=-1) |
| 86 | + |
| 87 | + # TODO: use torch._assert_async() when pytorch 1.8 support is dropped |
| 88 | + assert (x2 >= x1).all(), "bad box: x1 larger than x2" |
| 89 | + assert (y2 >= y1).all(), "bad box: y1 larger than y2" |
| 90 | + |
| 91 | + # Intersection keypoints |
| 92 | + xkis1 = torch.max(x1, x1g) |
| 93 | + ykis1 = torch.max(y1, y1g) |
| 94 | + xkis2 = torch.min(x2, x2g) |
| 95 | + ykis2 = torch.min(y2, y2g) |
| 96 | + |
| 97 | + intsct = torch.zeros_like(x1) |
| 98 | + mask = (ykis2 > ykis1) & (xkis2 > xkis1) |
| 99 | + intsct[mask] = (xkis2[mask] - xkis1[mask]) * (ykis2[mask] - ykis1[mask]) |
| 100 | + union = (x2 - x1) * (y2 - y1) + (x2g - x1g) * (y2g - y1g) - intsct + eps |
| 101 | + iou = intsct / union |
| 102 | + |
| 103 | + # smallest enclosing box |
| 104 | + xc1 = torch.min(x1, x1g) |
| 105 | + yc1 = torch.min(y1, y1g) |
| 106 | + xc2 = torch.max(x2, x2g) |
| 107 | + yc2 = torch.max(y2, y2g) |
| 108 | + diag_len = ((xc2 - xc1) ** 2) + ((yc2 - yc1) ** 2) + eps |
| 109 | + |
| 110 | + # centers of boxes |
| 111 | + x_p = (x2 + x1) / 2 |
| 112 | + y_p = (y2 + y1) / 2 |
| 113 | + x_g = (x1g + x2g) / 2 |
| 114 | + y_g = (y1g + y2g) / 2 |
| 115 | + distance = ((x_p - x_g) ** 2) + ((y_p - y_g) ** 2) |
| 116 | + |
| 117 | + # width and height of boxes |
| 118 | + w_pred = x2 - x1 |
| 119 | + h_pred = y2 - y1 |
| 120 | + w_gt = x2g - x1g |
| 121 | + h_gt = y2g - y1g |
| 122 | + v = (4 / (math.pi ** 2)) * torch.pow((torch.atan(w_gt / h_gt) - torch.atan(w_pred / h_pred)), 2) |
| 123 | + with torch.no_grad(): |
| 124 | + alpha = v / (1 - iou + v + eps) |
| 125 | + |
| 126 | + # Eqn. (10) |
| 127 | + loss = 1 - iou + (distance / diag_len) + alpha * v |
| 128 | + if reduction == "mean": |
| 129 | + loss = loss.mean() if loss.numel() > 0 else 0.0 * loss.sum() |
| 130 | + elif reduction == "sum": |
| 131 | + loss = loss.sum() |
| 132 | + |
| 133 | + return loss |
0 commit comments