深入解析CIoU损失函数原理图：从数学推导到PyTorch实现

1次阅读

没有评论

共计 2741 个字符，预计需要花费 7 分钟才能阅读完成。

CIoU（Complete Intersection over Union）损失函数是目标检测中用来优化边界框回归的重要工具。相比传统的 IoU，它增加了中心点距离惩罚和长宽比一致性惩罚，使得边界框的回归更加精确。

CIoU 的损失函数公式如下：

$$
\mathcal{L}_{CIoU} = 1 – IoU + \frac{\rho^2(b, b^{gt})}{c^2} + \alpha v
$$

其中：
– $\rho^2(b, b^{gt})$ 表示预测框 $b$ 和真实框 $b^{gt}$ 中心点的欧氏距离。
– $c$ 是最小包围框的对角线长度。
– $\alpha$ 是权重系数，$v$ 是长宽比一致性惩罚项。

中心点距离惩罚 ：通过最小化预测框和真实框的中心点距离，确保边界框的位置更加准确。
长宽比惩罚 ：通过惩罚预测框和真实框的长宽比差异，使得边界框的形状更加匹配。

以下是 CIoU 损失函数的完整 PyTorch 实现：

import torch
import math

class CIoULoss(torch.nn.Module):
    def __init__(self):
        super(CIoULoss, self).__init__()

    def _calculate_ciou(self, box1, box2):
        """
        box1: [x1, y1, x2, y2]
        box2: [x1, y1, x2, y2]
        """
        # 确保输入张量的形状正确
        assert box1.shape == box2.shape, "Input tensors must have the same shape"

        # 计算 IoU
        inter_x1 = torch.max(box1[..., 0], box2[..., 0])
        inter_y1 = torch.max(box1[..., 1], box2[..., 1])
        inter_x2 = torch.min(box1[..., 2], box2[..., 2])
        inter_y2 = torch.min(box1[..., 3], box2[..., 3])

        inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)

        area1 = (box1[..., 2] - box1[..., 0]) * (box1[..., 3] - box1[..., 1])
        area2 = (box2[..., 2] - box2[..., 0]) * (box2[..., 3] - box2[..., 1])

        union_area = area1 + area2 - inter_area
        iou = inter_area / (union_area + 1e-6)

        # 计算中心点距离
        center1 = torch.stack([(box1[..., 0] + box1[..., 2]) / 2, (box1[..., 1] + box1[..., 3]) / 2], dim=-1)
        center2 = torch.stack([(box2[..., 0] + box2[..., 2]) / 2, (box2[..., 1] + box2[..., 3]) / 2], dim=-1)

        center_dist = torch.sum((center1 - center2) ** 2, dim=-1)

        # 计算最小包围框的对角线长度
        enclose_x1 = torch.min(box1[..., 0], box2[..., 0])
        enclose_y1 = torch.min(box1[..., 1], box2[..., 1])
        enclose_x2 = torch.max(box1[..., 2], box2[..., 2])
        enclose_y2 = torch.max(box1[..., 3], box2[..., 3])

        enclose_dist = (enclose_x2 - enclose_x1) ** 2 + (enclose_y2 - enclose_y1) ** 2

        # 计算长宽比惩罚
        w1 = box1[..., 2] - box1[..., 0]
        h1 = box1[..., 3] - box1[..., 1]
        w2 = box2[..., 2] - box2[..., 0]
        h2 = box2[..., 3] - box2[..., 1]

        arctan = torch.atan(w2 / (h2 + 1e-6)) - torch.atan(w1 / (h1 + 1e-6))
        v = (4 / (math.pi ** 2)) * torch.pow(arctan, 2)

        alpha = v / (1 - iou + v + 1e-6)

        ciou = 1 - iou + center_dist / (enclose_dist + 1e-6) + alpha * v

        return ciou

    def forward(self, pred, target):
        """
        pred: [batch_size, 4]
        target: [batch_size, 4]
        """
        # 检查输入形状
        assert pred.shape == target.shape, "Pred and target must have the same shape"

        loss = self._calculate_ciou(pred, target)
        return loss.mean()