一、本文介绍

本文记录的是改进 RT-DETR 的损失函数，将其替换成 Shape-NWD 。 Shape-NWD 是一种用于目标检测中边界框回归的方法，它是在 Shape-IoU 和 NWD（Normalized Wasserstein Distance） 的基础上进行改进的，主要为了 更好地处理边界框回归中形状和尺度因素对回归结果的影响 ，特别是在 小目标检测任务 中。

---

---

二、Shape-NWD设计原理

Shape-NWD 是一种用于目标检测中边界框回归的方法，它是在 Shape-IoU 和 NWD（Normalized Wasserstein Distance） 的基础上进行改进的，主要为了 更好地处理边界框回归中形状和尺度因素对回归结果的影响 ，特别是在 小目标检测任务 中。

以下是关于 Shape-NWD 的详细介绍：

2.1 出发点

1. 现有方法的局限性 ：传统的基于 IoU 的边界框回归方法（如IoU、GIoU、CIoU等）在计算损失时，主要考虑预测框和真实框（GT box）之间的几何关系，如重叠面积、中心点距离等，但 忽略了边界框自身形状和尺度对回归结果的影响 。在 小目标检测场景下，这些方法对IoU值的敏感性问题更加突出 。
2. 小目标检测的需求 ： 小目标在图像中所占像素较少，其特征相对不明显，使得准确检测小目标变得更加困难 。因此，需要一种能够更精准地描述 小目标边界框回归损失 的方法，以提高小目标检测的性能。
3. 形状和尺度因素的重要性 ：通过对边界框回归特性的分析发现，回归样本中边界框的形状和尺度差异会导致在相同偏差情况下 IoU 值的不同。对于小尺度边界框，其形状和偏差对 IoU 值的影响更为显著。因此，在设计小目标检测的边界框回归损失函数时，有必要考虑 形状和尺度 因素。

2.2 原理

基于Shape-IoU的改进 ： Shape-NWD 将 Shape-IoU 的思想融入到 NWD 中。 Shape-IoU 在计算损失时引入了与 目标尺度相关的尺度因子 （scale）以及与GT box 形状相关的水平和垂直方向的权重系数 （ww和hh）。

2.2.1 NWD的基本形式

传统的 NWD 公式为 $$D=\sqrt{(x_c-x_c^{gt}{)}^2+(y_c-y_c^{gt}{)}^2+\frac{(w-w^{gt}{)}^2+(h-h^{gt}{)}^2}{weigh{t}^2}}$$ $$NWD=e^{-\frac{D}{C}}$$ 其中 $(x_c,y_c)$ 和 $(x_c^{gt},y_c^{gt})$ 分别是预测框和GT box的中心点坐标， $w$ 和 $h$ 是预测框的宽和高， $w^{gt}$ 和 $h^{gt}$ 是GT box的宽和高， $weight$ 是与数据集相关的常数， $C$ 也是与数据集相关的常数。

2.2.2 Shape-NWD的定义

$$B=\frac{(w-w_{gt}{)}^2+(h-h_{gt}{)}^2}{weigh{t}^2}$$ 其中 $weight=2$ 。
$$D=\sqrt{h_h\times (x_c-x_c^{gt}{)}^2+w_w\times (y_c-y_c^{gt}{)}^2+B}$$
$$NW{D}_{shape}=e^{-\frac{D}{C}}$$

• 公式中(D)是一个综合的距离度量， $x_c$ 和 $y_c$ 是预测框的中心点坐标， $x_c^{gt}$ 和 $y_c^{gt}$ 是真实框的中心点坐标， $h_h$ 和 $w_w$ 是与形状相关的系数， $w$ 和 $h$ 是预测框的宽和高， $w_{gt}$ 和 $h_{gt}$ 是真实框的宽和高）。
• $C$ 是一个与数据集相关的常数。

2.2.3 原理分析

1. 首先，计算 $D$ 时综合考虑了预测框和真实框在 中心点坐标差异 、 宽高差异 以及与 形状 相关的加权因素。 $(x_c-x_c^{gt}{)}^2$ 和 $(y_c-y_c^{gt}{)}^2$ 衡量了中心点的偏移，通过乘以 $h_h$ 和 $w_w$ ， 使得在水平和垂直方向上可以根据形状因素对这种偏移进行不同程度的加权 。例如，如果 $h_h$ 较大，那么在垂直方向上的中心点偏移对 $D$ 的贡献就会更大，这体现了对形状因素的考虑。
2. $B$ 部分计算了宽高的差异，同样通过除以 $weigh{t}^2$ （这里 $weight=2$ ）进行了归一化处理， 使得宽高差异在整个距离度量中具有合适的权重 。
3. 然后， $e^{-\frac{D}{C}}$ 这个指数形式的计算起到了 归一化 和 强调差异 的作用。当 $D$ 越大（即预测框和真实框的差异越大）时， $-\frac{D}{C}$ 的值越小， $e^{-\frac{D}{C}}$ 的值就越接近0，这表示预测框和真实框在形状方面的匹配程度越低；反之，当 $D$ 越小， $e^{-\frac{D}{C}}$ 的值越接近1，表示匹配程度越高。通过 $C$ 这个与数据集相关的常数，可以根据数据集的特点调整这种对差异的敏感度。

2.3 优势体现

• 形状和尺度敏感性 ：由于在计算 $D$ 时考虑了形状相关系数 $h_h$ 、 $w_w$ 以及宽高差异 $B$ ， Shape-NWD 能够有效捕捉边界框形状和尺度的变化。对于不同形状（如长条形和正方形）和尺度的边界框，它可以根据这些因素更准确地计算损失，避免了传统方法仅依赖于简单几何关系（如IoU）而忽略形状和尺度影响的问题。在小目标检测中，小目标的形状和尺度变化对检测结果影响较大， Shape-NWD 的这种敏感性能够更好地适应小目标的特性 ，提高小目标检测的准确性。
• 归一化处理 ： $e^{-\frac{D}{C}}$ 的形式对距离 $D$ 进行了归一化，使得 Shape-NWD 的值域在一定范围内（0到1之间），便于在损失函数计算和模型训练中进行比较和优化。这种归一化 有助于稳定模型训练过程 ， 避免因距离度量值过大或过小而导致的梯度问题 ，提高模型的收敛速度和稳定性。
• 适应性与泛化能力 ：通过引入与数据集相关的常数 $C$ ， Shape-NWD 可以根据不同数据集的特点进行调整。不同数据集可能包含具有不同形状和尺度分布的目标， $C$ 的值可以根据数据集的统计特性进行优化，从而使 Shape-NWD 在不同的目标检测任务中都能表现出较好的性能，具有较强的适应性和泛化能力。

论文： https://arxiv.org/pdf/2312.17663
源码： https://github.com/malagoutou/Shape-IoU

---

三、Shape-NWD的实现代码

Shape-NWD 的实现代码如下：

def shape_iou(box1, box2, xywh=True, scale=0, eps=1e-7):
    (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
    w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
    b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
    b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
 
    # Intersection area
    inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
            (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
 
    # Union Area
    union = w1 * h1 + w2 * h2 - inter + eps
 
    # IoU
    iou = inter / union
 
    #Shape-Distance    #Shape-Distance    #Shape-Distance    #Shape-Distance    #Shape-Distance    #Shape-Distance    #Shape-Distance  
    ww = 2 * torch.pow(w2, scale) / (torch.pow(w2, scale) + torch.pow(h2, scale))
    hh = 2 * torch.pow(h2, scale) / (torch.pow(w2, scale) + torch.pow(h2, scale))
    cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex width
    ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
    c2 = cw ** 2 + ch ** 2 + eps                            # convex diagonal squared
    center_distance_x = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2) / 4
    center_distance_y = ((b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4
    center_distance = hh * center_distance_x + ww * center_distance_y
    distance = center_distance / c2
  
    omiga_w = hh * torch.abs(w1 - w2) / torch.max(w1, w2)
    omiga_h = ww * torch.abs(h1 - h2) / torch.max(h1, h2)
    shape_cost = torch.pow(1 - torch.exp(-1 * omiga_w), 4) + torch.pow(1 - torch.exp(-1 * omiga_h), 4)
    
    iou = iou - distance - 0.5 * ( shape_cost)
    return iou  # IoU

def wasserstein_loss(pred, target, eps=1e-7, constant=12.8):
    r"""`Implementation of paper `Enhancing Geometric Factors into
    Model Learning and Inference for Object Detection and Instance
    Segmentation <https://arxiv.org/abs/2005.03572>`_.
    Code is modified from https://github.com/Zzh-tju/CIoU.
    Args:
        pred (Tensor): Predicted bboxes of format (x_center, y_center, w, h),
            shape (n, 4).
        target (Tensor): Corresponding gt bboxes, shape (n, 4).
        eps (float): Eps to avoid log(0).
    Return:
        Tensor: Loss tensor.
    """

    center1 = pred[:, :2]
    center2 = target[:, :2]

    whs = center1[:, :2] - center2[:, :2]

    center_distance = whs[:, 0] * whs[:, 0] + whs[:, 1] * whs[:, 1] + eps #

    w1 = pred[:, 2]  + eps
    h1 = pred[:, 3]  + eps
    w2 = target[:, 2] + eps
    h2 = target[:, 3] + eps

    wh_distance = ((w1 - w2) ** 2 + (h1 - h2) ** 2) / 4

    wasserstein_2 = center_distance + wh_distance
    return torch.exp(-torch.sqrt(wasserstein_2) / constant)

---

四、添加步骤

4.1 修改ultralytics/utils/metrics.py

此处需要查看的文件是 ultralytics/utils/metrics.py

metrics.py 中定义了模型的损失函数和计算方法，我们想要加入新的损失函数就只需要将代码放到这个文件内即可

Shape-NWD 添加后如下：

4.2 修改ultralytics/utils/loss.py

utils\loss.py 用于计算各种损失。

在 ultralytics/utils/loss.py 在的引用中添加 shape_iou , wasserstein_loss ，然后在 BboxLoss函数 内修改如下代码，使模型调用此 Shape-NWD 损失函数。

def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask):
        """IoU loss."""
        weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
        # iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
        iou = shape_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask])

        loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum

        nwd = wasserstein_loss(pred_bboxes[fg_mask], target_bboxes[fg_mask])
        iou_ratio = 0.3  #可调超参数，小目标可变低
        nwd_loss = ((1.0 - nwd) * weight).sum() / target_scores_sum
        loss_iou = iou_ratio * loss_iou + (1 - iou_ratio) * nwd_loss

4.3 修改ultralytics/utils/tal.py

tal.py 中是一些损失函数的功能应用。

在 ultralytics/utils/tal.py 在的引用中添加 shape_iou ，然后在 iou_calculation函数 内修改如下代码，使模型调用此 Shape-IoU 损失函数。

此时再次训练模型便会使用 Shape-NWD 计算模型的损失函数。

五、成功运行结果