|
- # Copyright (c) OpenMMLab. All rights reserved.
- import math
- import warnings
-
- import mmcv
- import torch
- import torch.nn as nn
-
- from mmdet.core import bbox_overlaps
- from ..builder import LOSSES
- from .utils import weighted_loss
-
-
- @mmcv.jit(derivate=True, coderize=True)
- @weighted_loss
- def iou_loss(pred, target, linear=False, mode='log', eps=1e-6):
- """IoU loss.
-
- Computing the IoU loss between a set of predicted bboxes and target bboxes.
- The loss is calculated as negative log of IoU.
-
- Args:
- pred (torch.Tensor): Predicted bboxes of format (x1, y1, x2, y2),
- shape (n, 4).
- target (torch.Tensor): Corresponding gt bboxes, shape (n, 4).
- linear (bool, optional): If True, use linear scale of loss instead of
- log scale. Default: False.
- mode (str): Loss scaling mode, including "linear", "square", and "log".
- Default: 'log'
- eps (float): Eps to avoid log(0).
-
- Return:
- torch.Tensor: Loss tensor.
- """
- assert mode in ['linear', 'square', 'log']
- if linear:
- mode = 'linear'
- warnings.warn('DeprecationWarning: Setting "linear=True" in '
- 'iou_loss is deprecated, please use "mode=`linear`" '
- 'instead.')
- ious = bbox_overlaps(pred, target, is_aligned=True).clamp(min=eps)
- if mode == 'linear':
- loss = 1 - ious
- elif mode == 'square':
- loss = 1 - ious**2
- elif mode == 'log':
- loss = -ious.log()
- else:
- raise NotImplementedError
- return loss
-
-
- @mmcv.jit(derivate=True, coderize=True)
- @weighted_loss
- def bounded_iou_loss(pred, target, beta=0.2, eps=1e-3):
- """BIoULoss.
-
- This is an implementation of paper
- `Improving Object Localization with Fitness NMS and Bounded IoU Loss.
- <https://arxiv.org/abs/1711.00164>`_.
-
- Args:
- pred (torch.Tensor): Predicted bboxes.
- target (torch.Tensor): Target bboxes.
- beta (float): beta parameter in smoothl1.
- eps (float): eps to avoid NaN.
- """
- pred_ctrx = (pred[:, 0] + pred[:, 2]) * 0.5
- pred_ctry = (pred[:, 1] + pred[:, 3]) * 0.5
- pred_w = pred[:, 2] - pred[:, 0]
- pred_h = pred[:, 3] - pred[:, 1]
- with torch.no_grad():
- target_ctrx = (target[:, 0] + target[:, 2]) * 0.5
- target_ctry = (target[:, 1] + target[:, 3]) * 0.5
- target_w = target[:, 2] - target[:, 0]
- target_h = target[:, 3] - target[:, 1]
-
- dx = target_ctrx - pred_ctrx
- dy = target_ctry - pred_ctry
-
- loss_dx = 1 - torch.max(
- (target_w - 2 * dx.abs()) /
- (target_w + 2 * dx.abs() + eps), torch.zeros_like(dx))
- loss_dy = 1 - torch.max(
- (target_h - 2 * dy.abs()) /
- (target_h + 2 * dy.abs() + eps), torch.zeros_like(dy))
- loss_dw = 1 - torch.min(target_w / (pred_w + eps), pred_w /
- (target_w + eps))
- loss_dh = 1 - torch.min(target_h / (pred_h + eps), pred_h /
- (target_h + eps))
- # view(..., -1) does not work for empty tensor
- loss_comb = torch.stack([loss_dx, loss_dy, loss_dw, loss_dh],
- dim=-1).flatten(1)
-
- loss = torch.where(loss_comb < beta, 0.5 * loss_comb * loss_comb / beta,
- loss_comb - 0.5 * beta)
- return loss
-
-
- @mmcv.jit(derivate=True, coderize=True)
- @weighted_loss
- def giou_loss(pred, target, eps=1e-7):
- r"""`Generalized Intersection over Union: A Metric and A Loss for Bounding
- Box Regression <https://arxiv.org/abs/1902.09630>`_.
-
- Args:
- pred (torch.Tensor): Predicted bboxes of format (x1, y1, x2, y2),
- shape (n, 4).
- target (torch.Tensor): Corresponding gt bboxes, shape (n, 4).
- eps (float): Eps to avoid log(0).
-
- Return:
- Tensor: Loss tensor.
- """
- gious = bbox_overlaps(pred, target, mode='giou', is_aligned=True, eps=eps)
- loss = 1 - gious
- return loss
-
-
- @mmcv.jit(derivate=True, coderize=True)
- @weighted_loss
- def diou_loss(pred, target, eps=1e-7):
- r"""`Implementation of Distance-IoU Loss: Faster and Better
- Learning for Bounding Box Regression, https://arxiv.org/abs/1911.08287`_.
-
- Code is modified from https://github.com/Zzh-tju/DIoU.
-
- Args:
- pred (Tensor): Predicted bboxes of format (x1, y1, x2, y2),
- shape (n, 4).
- target (Tensor): Corresponding gt bboxes, shape (n, 4).
- eps (float): Eps to avoid log(0).
- Return:
- Tensor: Loss tensor.
- """
- # overlap
- lt = torch.max(pred[:, :2], target[:, :2])
- rb = torch.min(pred[:, 2:], target[:, 2:])
- wh = (rb - lt).clamp(min=0)
- overlap = wh[:, 0] * wh[:, 1]
-
- # union
- ap = (pred[:, 2] - pred[:, 0]) * (pred[:, 3] - pred[:, 1])
- ag = (target[:, 2] - target[:, 0]) * (target[:, 3] - target[:, 1])
- union = ap + ag - overlap + eps
-
- # IoU
- ious = overlap / union
-
- # enclose area
- enclose_x1y1 = torch.min(pred[:, :2], target[:, :2])
- enclose_x2y2 = torch.max(pred[:, 2:], target[:, 2:])
- enclose_wh = (enclose_x2y2 - enclose_x1y1).clamp(min=0)
-
- cw = enclose_wh[:, 0]
- ch = enclose_wh[:, 1]
-
- c2 = cw**2 + ch**2 + eps
-
- b1_x1, b1_y1 = pred[:, 0], pred[:, 1]
- b1_x2, b1_y2 = pred[:, 2], pred[:, 3]
- b2_x1, b2_y1 = target[:, 0], target[:, 1]
- b2_x2, b2_y2 = target[:, 2], target[:, 3]
-
- left = ((b2_x1 + b2_x2) - (b1_x1 + b1_x2))**2 / 4
- right = ((b2_y1 + b2_y2) - (b1_y1 + b1_y2))**2 / 4
- rho2 = left + right
-
- # DIoU
- dious = ious - rho2 / c2
- loss = 1 - dious
- return loss
-
-
- @mmcv.jit(derivate=True, coderize=True)
- @weighted_loss
- def ciou_loss(pred, target, eps=1e-7):
- r"""`Implementation of paper `Enhancing Geometric Factors into
- Model Learning and Inference for Object Detection and Instance
- Segmentation <https://arxiv.org/abs/2005.03572>`_.
-
- Code is modified from https://github.com/Zzh-tju/CIoU.
-
- Args:
- pred (Tensor): Predicted bboxes of format (x1, y1, x2, y2),
- shape (n, 4).
- target (Tensor): Corresponding gt bboxes, shape (n, 4).
- eps (float): Eps to avoid log(0).
- Return:
- Tensor: Loss tensor.
- """
- # overlap
- lt = torch.max(pred[:, :2], target[:, :2])
- rb = torch.min(pred[:, 2:], target[:, 2:])
- wh = (rb - lt).clamp(min=0)
- overlap = wh[:, 0] * wh[:, 1]
-
- # union
- ap = (pred[:, 2] - pred[:, 0]) * (pred[:, 3] - pred[:, 1])
- ag = (target[:, 2] - target[:, 0]) * (target[:, 3] - target[:, 1])
- union = ap + ag - overlap + eps
-
- # IoU
- ious = overlap / union
-
- # enclose area
- enclose_x1y1 = torch.min(pred[:, :2], target[:, :2])
- enclose_x2y2 = torch.max(pred[:, 2:], target[:, 2:])
- enclose_wh = (enclose_x2y2 - enclose_x1y1).clamp(min=0)
-
- cw = enclose_wh[:, 0]
- ch = enclose_wh[:, 1]
-
- c2 = cw**2 + ch**2 + eps
-
- b1_x1, b1_y1 = pred[:, 0], pred[:, 1]
- b1_x2, b1_y2 = pred[:, 2], pred[:, 3]
- b2_x1, b2_y1 = target[:, 0], target[:, 1]
- b2_x2, b2_y2 = target[:, 2], target[:, 3]
-
- w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
- w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
-
- left = ((b2_x1 + b2_x2) - (b1_x1 + b1_x2))**2 / 4
- right = ((b2_y1 + b2_y2) - (b1_y1 + b1_y2))**2 / 4
- rho2 = left + right
-
- factor = 4 / math.pi**2
- v = factor * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
-
- with torch.no_grad():
- alpha = (ious > 0.5).float() * v / (1 - ious + v)
-
- # CIoU
- cious = ious - (rho2 / c2 + alpha * v)
- loss = 1 - cious.clamp(min=-1.0, max=1.0)
- return loss
-
-
- @LOSSES.register_module()
- class IoULoss(nn.Module):
- """IoULoss.
-
- Computing the IoU loss between a set of predicted bboxes and target bboxes.
-
- Args:
- linear (bool): If True, use linear scale of loss else determined
- by mode. Default: False.
- eps (float): Eps to avoid log(0).
- reduction (str): Options are "none", "mean" and "sum".
- loss_weight (float): Weight of loss.
- mode (str): Loss scaling mode, including "linear", "square", and "log".
- Default: 'log'
- """
-
- def __init__(self,
- linear=False,
- eps=1e-6,
- reduction='mean',
- loss_weight=1.0,
- mode='log'):
- super(IoULoss, self).__init__()
- assert mode in ['linear', 'square', 'log']
- if linear:
- mode = 'linear'
- warnings.warn('DeprecationWarning: Setting "linear=True" in '
- 'IOULoss is deprecated, please use "mode=`linear`" '
- 'instead.')
- self.mode = mode
- self.linear = linear
- self.eps = eps
- self.reduction = reduction
- self.loss_weight = loss_weight
-
- def forward(self,
- pred,
- target,
- weight=None,
- avg_factor=None,
- reduction_override=None,
- **kwargs):
- """Forward function.
-
- Args:
- pred (torch.Tensor): The prediction.
- target (torch.Tensor): The learning target of the prediction.
- weight (torch.Tensor, optional): The weight of loss for each
- prediction. Defaults to None.
- avg_factor (int, optional): Average factor that is used to average
- the loss. Defaults to None.
- reduction_override (str, optional): The reduction method used to
- override the original reduction method of the loss.
- Defaults to None. Options are "none", "mean" and "sum".
- """
- assert reduction_override in (None, 'none', 'mean', 'sum')
- reduction = (
- reduction_override if reduction_override else self.reduction)
- if (weight is not None) and (not torch.any(weight > 0)) and (
- reduction != 'none'):
- if pred.dim() == weight.dim() + 1:
- weight = weight.unsqueeze(1)
- return (pred * weight).sum() # 0
- if weight is not None and weight.dim() > 1:
- # TODO: remove this in the future
- # reduce the weight of shape (n, 4) to (n,) to match the
- # iou_loss of shape (n,)
- assert weight.shape == pred.shape
- weight = weight.mean(-1)
- loss = self.loss_weight * iou_loss(
- pred,
- target,
- weight,
- mode=self.mode,
- eps=self.eps,
- reduction=reduction,
- avg_factor=avg_factor,
- **kwargs)
- return loss
-
-
- @LOSSES.register_module()
- class BoundedIoULoss(nn.Module):
-
- def __init__(self, beta=0.2, eps=1e-3, reduction='mean', loss_weight=1.0):
- super(BoundedIoULoss, self).__init__()
- self.beta = beta
- self.eps = eps
- self.reduction = reduction
- self.loss_weight = loss_weight
-
- def forward(self,
- pred,
- target,
- weight=None,
- avg_factor=None,
- reduction_override=None,
- **kwargs):
- if weight is not None and not torch.any(weight > 0):
- if pred.dim() == weight.dim() + 1:
- weight = weight.unsqueeze(1)
- return (pred * weight).sum() # 0
- assert reduction_override in (None, 'none', 'mean', 'sum')
- reduction = (
- reduction_override if reduction_override else self.reduction)
- loss = self.loss_weight * bounded_iou_loss(
- pred,
- target,
- weight,
- beta=self.beta,
- eps=self.eps,
- reduction=reduction,
- avg_factor=avg_factor,
- **kwargs)
- return loss
-
-
- @LOSSES.register_module()
- class GIoULoss(nn.Module):
-
- def __init__(self, eps=1e-6, reduction='mean', loss_weight=1.0):
- super(GIoULoss, self).__init__()
- self.eps = eps
- self.reduction = reduction
- self.loss_weight = loss_weight
-
- def forward(self,
- pred,
- target,
- weight=None,
- avg_factor=None,
- reduction_override=None,
- **kwargs):
- if weight is not None and not torch.any(weight > 0):
- if pred.dim() == weight.dim() + 1:
- weight = weight.unsqueeze(1)
- return (pred * weight).sum() # 0
- assert reduction_override in (None, 'none', 'mean', 'sum')
- reduction = (
- reduction_override if reduction_override else self.reduction)
- if weight is not None and weight.dim() > 1:
- # TODO: remove this in the future
- # reduce the weight of shape (n, 4) to (n,) to match the
- # giou_loss of shape (n,)
- assert weight.shape == pred.shape
- weight = weight.mean(-1)
- loss = self.loss_weight * giou_loss(
- pred,
- target,
- weight,
- eps=self.eps,
- reduction=reduction,
- avg_factor=avg_factor,
- **kwargs)
- return loss
-
-
- @LOSSES.register_module()
- class DIoULoss(nn.Module):
-
- def __init__(self, eps=1e-6, reduction='mean', loss_weight=1.0):
- super(DIoULoss, self).__init__()
- self.eps = eps
- self.reduction = reduction
- self.loss_weight = loss_weight
-
- def forward(self,
- pred,
- target,
- weight=None,
- avg_factor=None,
- reduction_override=None,
- **kwargs):
- if weight is not None and not torch.any(weight > 0):
- if pred.dim() == weight.dim() + 1:
- weight = weight.unsqueeze(1)
- return (pred * weight).sum() # 0
- assert reduction_override in (None, 'none', 'mean', 'sum')
- reduction = (
- reduction_override if reduction_override else self.reduction)
- if weight is not None and weight.dim() > 1:
- # TODO: remove this in the future
- # reduce the weight of shape (n, 4) to (n,) to match the
- # giou_loss of shape (n,)
- assert weight.shape == pred.shape
- weight = weight.mean(-1)
- loss = self.loss_weight * diou_loss(
- pred,
- target,
- weight,
- eps=self.eps,
- reduction=reduction,
- avg_factor=avg_factor,
- **kwargs)
- return loss
-
-
- @LOSSES.register_module()
- class CIoULoss(nn.Module):
-
- def __init__(self, eps=1e-6, reduction='mean', loss_weight=1.0):
- super(CIoULoss, self).__init__()
- self.eps = eps
- self.reduction = reduction
- self.loss_weight = loss_weight
-
- def forward(self,
- pred,
- target,
- weight=None,
- avg_factor=None,
- reduction_override=None,
- **kwargs):
- if weight is not None and not torch.any(weight > 0):
- if pred.dim() == weight.dim() + 1:
- weight = weight.unsqueeze(1)
- return (pred * weight).sum() # 0
- assert reduction_override in (None, 'none', 'mean', 'sum')
- reduction = (
- reduction_override if reduction_override else self.reduction)
- if weight is not None and weight.dim() > 1:
- # TODO: remove this in the future
- # reduce the weight of shape (n, 4) to (n,) to match the
- # giou_loss of shape (n,)
- assert weight.shape == pred.shape
- weight = weight.mean(-1)
- loss = self.loss_weight * ciou_loss(
- pred,
- target,
- weight,
- eps=self.eps,
- reduction=reduction,
- avg_factor=avg_factor,
- **kwargs)
- return loss
|