File size: 4,247 Bytes
34d1f8b |
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 116 117 118 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional
import torch
from mmdet.models.losses.utils import weighted_loss
from torch import Tensor
from torch import nn as nn
from mmdet3d.models import axis_aligned_iou_loss
from mmdet3d.registry import MODELS
from mmdet3d.structures import AxisAlignedBboxOverlaps3D
@weighted_loss
def axis_aligned_diou_loss(pred: Tensor, target: Tensor) -> Tensor:
"""Calculate the DIoU loss (1-DIoU) of two sets of axis aligned bounding
boxes. Note that predictions and targets are one-to-one corresponded.
Args:
pred (torch.Tensor): Bbox predictions with shape [..., 6]
(x1, y1, z1, x2, y2, z2).
target (torch.Tensor): Bbox targets (gt) with shape [..., 6]
(x1, y1, z1, x2, y2, z2).
Returns:
torch.Tensor: DIoU loss between predictions and targets.
"""
axis_aligned_iou = AxisAlignedBboxOverlaps3D()(
pred, target, is_aligned=True)
iou_loss = 1 - axis_aligned_iou
xp1, yp1, zp1, xp2, yp2, zp2 = pred.split(1, dim=-1)
xt1, yt1, zt1, xt2, yt2, zt2 = target.split(1, dim=-1)
xpc = (xp1 + xp2) / 2
ypc = (yp1 + yp2) / 2
zpc = (zp1 + zp2) / 2
xtc = (xt1 + xt2) / 2
ytc = (yt1 + yt2) / 2
ztc = (zt1 + zt2) / 2
r2 = (xpc - xtc)**2 + (ypc - ytc)**2 + (zpc - ztc)**2
x_min = torch.minimum(xp1, xt1)
x_max = torch.maximum(xp2, xt2)
y_min = torch.minimum(yp1, yt1)
y_max = torch.maximum(yp2, yt2)
z_min = torch.minimum(zp1, zt1)
z_max = torch.maximum(zp2, zt2)
c2 = (x_min - x_max)**2 + (y_min - y_max)**2 + (z_min - z_max)**2
diou_loss = iou_loss + (r2 / c2)[:, 0]
return diou_loss
@MODELS.register_module()
class TR3DAxisAlignedIoULoss(nn.Module):
"""Calculate the IoU loss (1-IoU) of axis aligned bounding boxes. The only
difference with original AxisAlignedIoULoss is the addition of DIoU mode.
These classes should be merged in the future.
Args:
mode (str): 'iou' for intersection over union or 'diou' for
distance-iou loss. Defaults to 'iou'.
reduction (str): Method to reduce losses.
The valid reduction method are 'none', 'sum' or 'mean'.
Defaults to 'mean'.
loss_weight (float): Weight of loss. Defaults to 1.0.
"""
def __init__(self,
mode: str = 'iou',
reduction: str = 'mean',
loss_weight: float = 1.0) -> None:
super(TR3DAxisAlignedIoULoss, self).__init__()
assert mode in ['iou', 'diou']
self.loss = axis_aligned_iou_loss if mode == 'iou' \
else axis_aligned_diou_loss
assert reduction in ['none', 'sum', 'mean']
self.reduction = reduction
self.loss_weight = loss_weight
def forward(self,
pred: Tensor,
target: Tensor,
weight: Optional[Tensor] = None,
avg_factor: Optional[float] = None,
reduction_override: Optional[str] = None,
**kwargs) -> Tensor:
"""Forward function of loss calculation.
Args:
pred (Tensor): Bbox predictions with shape [..., 3].
target (Tensor): Bbox targets (gt) with shape [..., 3].
weight (Tensor, optional): Weight of loss.
Defaults to None.
avg_factor (float, optional): Average factor that is used to
average the loss. Defaults to None.
reduction_override (str, optional): Method to reduce losses.
The valid reduction method are 'none', 'sum' or 'mean'.
Defaults to None.
Returns:
Tensor: IoU loss between predictions and targets.
"""
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'):
return (pred * weight).sum()
return self.loss(
pred,
target,
weight=weight,
avg_factor=avg_factor,
reduction=reduction) * self.loss_weight
|