Spaces:
Build error
Build error
File size: 4,982 Bytes
6d7be64 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import math
import torch
class IOUloss:
""" Calculate IoU loss.
"""
def __init__(self, box_format='xywh', iou_type='ciou', reduction='none', eps=1e-7):
""" Setting of the class.
Args:
box_format: (string), must be one of 'xywh' or 'xyxy'.
iou_type: (string), can be one of 'ciou', 'diou', 'giou' or 'siou'
reduction: (string), specifies the reduction to apply to the output, must be one of 'none', 'mean','sum'.
eps: (float), a value to avoid divide by zero error.
"""
self.box_format = box_format
self.iou_type = iou_type.lower()
self.reduction = reduction
self.eps = eps
def __call__(self, box1, box2):
""" calculate iou. box1 and box2 are torch tensor with shape [M, 4] and [Nm 4].
"""
box2 = box2.T
if self.box_format == 'xyxy':
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
elif self.box_format == 'xywh':
b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
# Intersection area
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
# Union Area
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + self.eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + self.eps
union = w1 * h1 + w2 * h2 - inter + self.eps
iou = inter / union
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex width
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
if self.iou_type == 'giou':
c_area = cw * ch + self.eps # convex area
iou = iou - (c_area - union) / c_area
elif self.iou_type in ['diou', 'ciou']:
c2 = cw ** 2 + ch ** 2 + self.eps # convex diagonal squared
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
(b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared
if self.iou_type == 'diou':
iou = iou - rho2 / c2
elif self.iou_type == 'ciou':
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
with torch.no_grad():
alpha = v / (v - iou + (1 + self.eps))
iou = iou - (rho2 / c2 + v * alpha)
elif self.iou_type == 'siou':
# SIoU Loss https://arxiv.org/pdf/2205.12740.pdf
s_cw = (b2_x1 + b2_x2 - b1_x1 - b1_x2) * 0.5
s_ch = (b2_y1 + b2_y2 - b1_y1 - b1_y2) * 0.5
sigma = torch.pow(s_cw ** 2 + s_ch ** 2, 0.5)
sin_alpha_1 = torch.abs(s_cw) / sigma
sin_alpha_2 = torch.abs(s_ch) / sigma
threshold = pow(2, 0.5) / 2
sin_alpha = torch.where(sin_alpha_1 > threshold, sin_alpha_2, sin_alpha_1)
angle_cost = torch.cos(torch.arcsin(sin_alpha) * 2 - math.pi / 2)
rho_x = (s_cw / cw) ** 2
rho_y = (s_ch / ch) ** 2
gamma = angle_cost - 2
distance_cost = 2 - torch.exp(gamma * rho_x) - torch.exp(gamma * rho_y)
omiga_w = torch.abs(w1 - w2) / torch.max(w1, w2)
omiga_h = torch.abs(h1 - h2) / torch.max(h1, h2)
shape_cost = torch.pow(1 - torch.exp(-1 * omiga_w), 4) + torch.pow(1 - torch.exp(-1 * omiga_h), 4)
iou = iou - 0.5 * (distance_cost + shape_cost)
loss = 1.0 - iou
if self.reduction == 'sum':
loss = loss.sum()
elif self.reduction == 'mean':
loss = loss.mean()
return loss
def pairwise_bbox_iou(box1, box2, box_format='xywh'):
"""Calculate iou.
This code is based on https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/utils/boxes.py
"""
if box_format == 'xyxy':
lt = torch.max(box1[:, None, :2], box2[:, :2])
rb = torch.min(box1[:, None, 2:], box2[:, 2:])
area_1 = torch.prod(box1[:, 2:] - box1[:, :2], 1)
area_2 = torch.prod(box2[:, 2:] - box2[:, :2], 1)
elif box_format == 'xywh':
lt = torch.max(
(box1[:, None, :2] - box1[:, None, 2:] / 2),
(box2[:, :2] - box2[:, 2:] / 2),
)
rb = torch.min(
(box1[:, None, :2] + box1[:, None, 2:] / 2),
(box2[:, :2] + box2[:, 2:] / 2),
)
area_1 = torch.prod(box1[:, 2:], 1)
area_2 = torch.prod(box2[:, 2:], 1)
valid = (lt < rb).type(lt.type()).prod(dim=2)
inter = torch.prod(rb - lt, 2) * valid
return inter / (area_1[:, None] + area_2 - inter)
|