IOU、GIOU、DIOU、CIOU 代码实现

废话不多说,原理看这个:

IOU、GIOU、DIOU、CIOU详解及代码实现

代码如下

为了更好的处理多个bbox的IoU计算，对代码进行了改造

import math
import torch
from torch import Tensor
# author:wuliang
# time: 2022.5.27

def _upcast(t: Tensor) -> Tensor:
    # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
    if t.is_floating_point():
        return t if t.dtype in (torch.float32, torch.float64) else t.float()
    else:
        return t if t.dtype in (torch.int32, torch.int64) else t.int()

def box_area(boxes: Tensor) -> Tensor:
    """
    Computes the area of a set of bounding boxes, which are specified by their
    (x1, y1, x2, y2) coordinates.

    Args:
        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
            are expected to be in (x1, y1, x2, y2) format with
            ``0 <= x1 < x2`` and ``0 <= y1 < y2``.

    Returns:
        Tensor[N]: the area for each box
    """
    boxes = _upcast(boxes)
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union


def all_bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
    """在ComputeLoss的__call__函数中调用计算回归损失
    :params box1: gt框  [nums1,4]
    :params box2: 预测框 [nums2,4]
    :return box1和box2的IoU/GIoU/DIoU/CIoU
    """
    w1= box1[:,2] -box1[:,0]
    h1= box1[:,3] -box1[:,1]+ eps

    w2 = box2[:, 2] - box2[:, 0]
    h2 = box2[:, 3] - box2[:, 1] + eps

    # iou = inter / union
    iou, union = box_iou(box1, box2)  # 200 12

    lt = torch.min(box1[:, None, :2], box2[:, :2])  # 200 12 2
    rb = torch.max(box1[:, None, 2:], box2[:, 2:])  # 200 12 2

    if GIoU or DIoU or CIoU:
        wh = (rb - lt).clamp(min=0)  # [N,M,2]
        cw = wh[:, :, 0] # 两个框的最小闭包区域的width
        ch = wh[:, :, 1] # 两个框的最小闭包区域的height
        if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            part1 = box1[:, None, :2] - box2[:, :2]
            part2 = box1[:, None, 2:] - box2[:, 2:]
            rho2=(  (part1[:,:,0]+part1[:,:,1]) ** 2 +(part2[:,:,0]+part2[:,:,1]) )/4

            if DIoU:
                return iou - rho2 / c2  # DIoU
            elif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47

                v = (4 / math.pi ** 2) * torch.pow( torch.atan(w1 / h1)[:,None] -torch.atan(w2 / h2) , 2)

                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU
        else:  # GIoU https://arxiv.org/pdf/1902.09630.pdf
            c_area = cw * ch + eps  # convex area
            return iou - (c_area - union) / c_area  # GIoU
    else:
        return iou  # IoU


if __name__ == '__main__':
    nums=100
	# 方便测试，bbox1，和bbox2 设置了固定值
    bbox1=torch.ones((2,4))
    init1=torch.tensor([1,2,3,4])[None,:]

    bbox2 = torch.ones((1,4))
    init2 = torch.tensor([2, 3, 4, 5])[None,:]

    bbox1 =bbox1*init1
    bbox2 =bbox2*init2

    # iou,_  =box_iou(bbox1, bbox2)
    iou = all_bbox_iou(bbox1, bbox2, )
    costiou = -iou
    print('iou', costiou[0])

    giou = all_bbox_iou(bbox1, bbox2,GIoU=True)
    costgiou = -giou
    print('giou', costgiou[0])

    diou = all_bbox_iou(bbox1,bbox2,DIoU=True)
    costdiou =  -diou
    print('diou',costdiou[0])

    ciou = all_bbox_iou(bbox1, bbox2, CIoU=True)
    costciou = -ciou
    print('ciou',costciou[0])