giou loss

https://github.com/Tenvence/gravel-object-detection/blob/0feea9b29bbbd0a6cc02516e7ecffc27ef732bf7/model/loss.py

import numpy as np
import torch
import torch.nn as nn

def compute_giou(boxes1, boxes2):
    boxes1 = regularize_boxes(boxes1)
    boxes2 = regularize_boxes(boxes2)

    inter_area, union_area, enclose_area = compute_inter_union_area(boxes1, boxes2)

    iou = inter_area / union_area
    giou = iou - (enclose_area - union_area) / enclose_area

    return giou


def regularize_boxes(boxes):
    boxes = torch.cat([boxes[..., :2] - boxes[..., 2:] * 0.5, boxes[..., :2] + boxes[..., 2:] * 0.5], dim=-1)
    boxes = torch.cat([torch.min(boxes[..., :2], boxes[..., 2:]), torch.max(boxes[..., :2], boxes[..., 2:])], dim=-1)

    return boxes


def compute_inter_union_area(boxes1, boxes2):
    #x1 y1 x2 y2
    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    left_up = torch.max(boxes1[..., :2], boxes2[..., :2])
    right_down = torch.min(boxes1[..., 2:], boxes2[..., 2:])

    zero = torch.zeros((1, 2))

    # Support GPU computation
    if boxes1.is_cuda:
        zero = zero.to(device=boxes1.device)

    inter_section = torch.max(right_down - left_up, zero)

    inter_area = inter_section[..., 0] * inter_section[..., 1]
    union_area = boxes1_area + boxes2_area - inter_area

    enclose_left_up = torch.min(boxes1[..., :2], boxes2[..., :2])
    enclose_right_down = torch.max(boxes1[..., 2:], boxes2[..., 2:])
    enclose = torch.max(enclose_right_down - enclose_left_up, zero)
    enclose_area = enclose[..., 0] * enclose[..., 1]

    return inter_area, union_area, enclose_area

class Loss(nn.Module):
    def __init__(self, anchors, input_size):
        super(Loss, self).__init__()

        self.s_anchors, self.m_anchors, self.l_anchors = anchors

        self.input_size = torch.tensor(input_size)
        self.iou_thresh = torch.tensor(0.5)

    def forward(self, s_output, m_output, l_output, s_gt_tensor, m_gt_tensor, l_gt_tensor, s_gt_coords, m_gt_coords, l_gt_coords):
        s_giou_loss = self.compute_loss(s_output, s_gt_tensor, s_gt_coords)
        m_giou_loss = self.compute_loss(m_output, m_gt_tensor, m_gt_coords)
        l_giou_loss = self.compute_loss(l_output, l_gt_tensor, l_gt_coords)


        return s_giou_loss + m_giou_loss + l_giou_loss

    def compute_loss(self, output, gt_tensor):
        output_coord = output[:, :, :, :, 0:4]
        output_conf = output[:, :, :, :, 4:5]

        gt_tensor_coord = gt_tensor[:, :, :, :, 0:4]
        gt_tensor_conf = gt_tensor[:, :, :, :, 4:5]

        giou_loss = self.compute_giou_loss(output_coord, gt_tensor_coord, gt_tensor_conf)
        # conf_loss = self.compute_conf_loss(output_coord, output_conf, gt_tensor_conf, gt_coords)

        giou_loss = torch.mean(torch.sum(giou_loss, [1, 2, 3, 4]))
        # conf_loss = torch.mean(torch.sum(conf_loss, [1, 2, 3, 4]))

        return giou_loss

    def compute_giou_loss(self, output_coord, gt_tensor_coord, gt_tensor_conf):
        giou = compute_giou(output_coord, gt_tensor_coord)[..., np.newaxis]
        box_scale = torch.div(gt_tensor_coord[:, :, :, :, 2:3] * gt_tensor_coord[:, :, :, :, 3:4], torch.pow(self.input_size, 2))
        giou_loss = gt_tensor_conf * torch.sub(2, box_scale) * torch.sub(1, giou)

        return giou_loss