目标检测中各种IoU说明

Posted 2023-03-02 Goodness2020

一、说明

IoU损失是目标检测中最常见的损失函数，表示的就是真实框和预测框的交并比。

二、代码实现

import math

def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
    # Returns Intersection over Union IoU(n,1) of box1(n,4) to box2(n,4)
    # Get the coordinates of bounding boxes
    if xywh:  # transform from xywh to xyxy
        x1, y1, w1, h1 = box1.chunk(4, -1)  # n*4 --> n*1, n*1, n*1, n*1
        x2, y2, w2, h2 = box2.chunk(4, -1)
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1 / 2, x1 + w1 / 2, y1 - h1 / 2, y1 + h1 / 2
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2 / 2, x2 + w2 / 2, y2 - h2 / 2, y2 + h2 / 2
    else:  # x1, y1, x2, y2 = box
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)  # n*4 --> n*1, n*1, n*1, n*1
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
        w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps

    # Intersection Area   n*1
    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * \\
            (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp(0)

    # Union Area   n*1
    union = w1 * h1 + w2 * h2 - inter + eps

    # IoU  n*1
    iou = inter / union
    if CIoU or DIoU or GIoU:
        cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # convex width
        ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # convex height
        if CIoU or DIoU:  # Distance or Complete IoU
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center dist ** 2
            if CIoU:
                v = (4 / math.pi ** 2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU = IoU - center_dist / diagonal_dist - length_width_ratio
            return iou - rho2 / c2  # DIoU = IoU - center_dist / diagonal_dist
        c_area = cw * ch + eps  # convex area
        return iou - (c_area - union) / c_area  # GIoU = IoU - (area_convex - area_union) / area_convex
    return iou  # IoU