| 8 | |
| 9 | |
| 10 | class IOUloss(nn.Module): |
| 11 | def __init__(self, reduction="none", loss_type="iou"): |
| 12 | super(IOUloss, self).__init__() |
| 13 | self.reduction = reduction |
| 14 | self.loss_type = loss_type |
| 15 | |
| 16 | def forward(self, pred, target): |
| 17 | assert pred.shape[0] == target.shape[0] |
| 18 | |
| 19 | pred = pred.view(-1, 4) |
| 20 | target = target.view(-1, 4) |
| 21 | tl = torch.max( |
| 22 | (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2) |
| 23 | ) |
| 24 | br = torch.min( |
| 25 | (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2) |
| 26 | ) |
| 27 | |
| 28 | area_p = torch.prod(pred[:, 2:], 1) |
| 29 | area_g = torch.prod(target[:, 2:], 1) |
| 30 | |
| 31 | en = (tl < br).type(tl.type()).prod(dim=1) |
| 32 | area_i = torch.prod(br - tl, 1) * en |
| 33 | iou = (area_i) / (area_p + area_g - area_i + 1e-16) |
| 34 | |
| 35 | if self.loss_type == "iou": |
| 36 | loss = 1 - iou ** 2 |
| 37 | elif self.loss_type == "giou": |
| 38 | c_tl = torch.min( |
| 39 | (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2) |
| 40 | ) |
| 41 | c_br = torch.max( |
| 42 | (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2) |
| 43 | ) |
| 44 | area_c = torch.prod(c_br - c_tl, 1) |
| 45 | giou = iou - (area_c - area_i) / area_c.clamp(1e-16) |
| 46 | loss = 1 - giou.clamp(min=-1.0, max=1.0) |
| 47 | |
| 48 | if self.reduction == "mean": |
| 49 | loss = loss.mean() |
| 50 | elif self.reduction == "sum": |
| 51 | loss = loss.sum() |
| 52 | |
| 53 | return loss |
| 54 | |
| 55 | |
| 56 | def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2): |