iou.py

import math

import torch
from torch import nn


class IouLoss(nn.Module):
    ''' :param monotonous: {
            None: origin
            True: monotonic FM
            False: non-monotonic FM
        }'''
    momentum = 1e-2
    alpha = 1.7
    delta = 2.7

    def __init__(self, ltype='WIoU', monotonous=False):
        super().__init__()
        assert getattr(self, f'_{ltype}', None), f'The loss function {ltype} does not exist'
        self.ltype = ltype
        self.monotonous = monotonous
        self.register_buffer('iou_mean', torch.tensor(1.))

    def __getitem__(self, item):
        if callable(self._fget[item]):
            self._fget[item] = self._fget[item]()
        return self._fget[item]

    def forward(self, pred, target, ret_iou=False, **kwargs):
        self._fget = {
            # pred, target: x0,y0,x1,y1
            'pred': pred,
            'target': target,
            # x,y,w,h
            'pred_xy': lambda: (self['pred'][..., :2] + self['pred'][..., 2: 4]) / 2,
            'pred_wh': lambda: self['pred'][..., 2: 4] - self['pred'][..., :2],
            'target_xy': lambda: (self['target'][..., :2] + self['target'][..., 2: 4]) / 2,
            'target_wh': lambda: self['target'][..., 2: 4] - self['target'][..., :2],
            # x0,y0,x1,y1
            'min_coord': lambda: torch.minimum(self['pred'][..., :4], self['target'][..., :4]),
            'max_coord': lambda: torch.maximum(self['pred'][..., :4], self['target'][..., :4]),
            # The overlapping region
            'wh_inter': lambda: torch.relu(self['min_coord'][..., 2: 4] - self['max_coord'][..., :2]),
            's_inter': lambda: torch.prod(self['wh_inter'], dim=-1),
            # The area covered
            's_union': lambda: torch.prod(self['pred_wh'], dim=-1) +
                               torch.prod(self['target_wh'], dim=-1) - self['s_inter'],
            # The smallest enclosing box
            'wh_box': lambda: self['max_coord'][..., 2: 4] - self['min_coord'][..., :2],
            's_box': lambda: torch.prod(self['wh_box'], dim=-1),
            'l2_box': lambda: torch.square(self['wh_box']).sum(dim=-1),
            # The central points' connection of the bounding boxes
            'd_center': lambda: self['pred_xy'] - self['target_xy'],
            'l2_center': lambda: torch.square(self['d_center']).sum(dim=-1),
            # IoU
            'iou': lambda: 1 - self['s_inter'] / self['s_union']
        }

        if self.training:
            self.iou_mean.mul_(1 - self.momentum)
            self.iou_mean.add_(self.momentum * self['iou'].detach().mean())

        ret = self._scaled_loss(getattr(self, f'_{self.ltype}')(**kwargs)), self['iou']
        delattr(self, '_fget')
        return ret if ret_iou else ret[0]

    def _scaled_loss(self, loss, iou=None):
        if isinstance(self.monotonous, bool):
            beta = (self['iou'].detach() if iou is None else iou) / self.iou_mean

            if self.monotonous:
                loss *= beta.sqrt()
            else:
                divisor = self.delta * torch.pow(self.alpha, beta - self.delta)
                loss *= beta / divisor
        return loss

    def _IoU(self):
        return self['iou']

    def _WIoU(self):
        dist = torch.exp(self['l2_center'] / self['l2_box'].detach())
        return dist * self['iou']

    def _EIoU(self):
        penalty = self['l2_center'] / self['l2_box'] \
                  + torch.square(self['d_center'] / self['wh_box']).sum(dim=-1)
        return self['iou'] + penalty

    def _GIoU(self):
        return self['iou'] + (self['s_box'] - self['s_union']) / self['s_box']

    def _DIoU(self):
        return self['iou'] + self['l2_center'] / self['l2_box']

    def _CIoU(self, eps=1e-4):
        v = 4 / math.pi ** 2 * \
            (torch.atan(self['pred_wh'][..., 0] / (self['pred_wh'][..., 1] + eps)) -
             torch.atan(self['target_wh'][..., 0] / (self['target_wh'][..., 1] + eps))) ** 2
        alpha = v / (self['iou'] + v)
        return self['iou'] + self['l2_center'] / self['l2_box'] + alpha.detach() * v

    def _SIoU(self, theta=4):
        # Angle Cost
        angle = torch.arcsin(torch.abs(self['d_center']).min(dim=-1)[0] / (self['l2_center'].sqrt() + 1e-4))
        angle = torch.sin(2 * angle) - 2
        # Dist Cost
        dist = angle[..., None] * torch.square(self['d_center'] / self['wh_box'])
        dist = 2 - torch.exp(dist[..., 0]) - torch.exp(dist[..., 1])
        # Shape Cost
        d_shape = torch.abs(self['pred_wh'] - self['target_wh'])
        big_shape = torch.maximum(self['pred_wh'], self['target_wh'])
        w_shape = 1 - torch.exp(- d_shape[..., 0] / big_shape[..., 0])
        h_shape = 1 - torch.exp(- d_shape[..., 1] / big_shape[..., 1])
        shape = w_shape ** theta + h_shape ** theta
        return self['iou'] + (dist + shape) / 2

    def __repr__(self):
        return f'{self.__name__}(iou_mean={self.iou_mean.item():.3f})'

    __name__ = property(lambda self: self.ltype)


if __name__ == '__main__':
    def xywh2xyxy(labels, i=0):
        labels = labels.clone()
        labels[..., i:i + 2] -= labels[..., i + 2:i + 4] / 2
        labels[..., i + 2:i + 4] += labels[..., i:i + 2]
        return labels


    torch.manual_seed(0)
    iouloss = IouLoss(ltype='WIoU').cuda()
    print(iouloss)

    for i in range(5):
        origin = torch.rand([2, 3, 1, 4], requires_grad=True, device=iouloss.iou_mean.device)
        pred, tar = xywh2xyxy(origin)

        loss = iouloss(pred, tar)
        loss.sum().backward()
        print(origin.grad)
    print(iouloss)