Hugging Face's logo

LightChuan
/
goat

Model card Files
xet
 Community
goat
File size: 4,774 Bytes
6a5bb7e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""Inner-IoU / Focal-IoU — 针对小目标定位精度的损失函数

Inner-IoU: 用缩放后的内部框计算IoU，让模型更关注框的内部重叠
           对小目标特别有效（60.8%的目标面积<0.5%）
Focal-IoU: 对高IoU样本加大权重，推动模型精修框的位置

论文: Inner-IoU: More Effective Intersection over Union Loss (2023)

Patch方式: 替换 BboxLoss.forward，不改模型结构
"""
import math
import torch
import torch.nn.functional as F


def inner_bbox_iou(box1, box2, xywh=False, ratio=0.7, eps=1e-7):
    """Inner-IoU: 用缩放后的内部框计算CIoU。

    Args:
        box1, box2: 预测框和目标框 (xyxy格式)
        ratio: 内部框缩放比例 (0.5~0.9, 越小越关注中心)
        eps: 防止除零
    """
    if xywh:
        (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_
    else:
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -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

    # Inner boxes: 以中心为基准缩放
    cx1, cy1 = (b1_x1 + b1_x2) / 2, (b1_y1 + b1_y2) / 2
    cx2, cy2 = (b2_x1 + b2_x2) / 2, (b2_y1 + b2_y2) / 2

    inner_w1, inner_h1 = w1 * ratio, h1 * ratio
    inner_w2, inner_h2 = w2 * ratio, h2 * ratio

    inner_b1_x1 = cx1 - inner_w1 / 2
    inner_b1_x2 = cx1 + inner_w1 / 2
    inner_b1_y1 = cy1 - inner_h1 / 2
    inner_b1_y2 = cy1 + inner_h1 / 2

    inner_b2_x1 = cx2 - inner_w2 / 2
    inner_b2_x2 = cx2 + inner_w2 / 2
    inner_b2_y1 = cy2 - inner_h2 / 2
    inner_b2_y2 = cy2 + inner_h2 / 2

    # Inner intersection
    inner_inter = (
        (inner_b1_x2.minimum(inner_b2_x2) - inner_b1_x1.maximum(inner_b2_x1)).clamp_(0)
        * (inner_b1_y2.minimum(inner_b2_y2) - inner_b1_y1.maximum(inner_b2_y1)).clamp_(0)
    )

    # Inner union
    inner_union = inner_w1 * inner_h1 + inner_w2 * inner_h2 - inner_inter + eps

    # Inner IoU
    inner_iou = inner_inter / inner_union

    # CIoU penalties (用原始框计算)
    cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)
    ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)
    c2 = cw.pow(2) + ch.pow(2) + eps
    rho2 = (
        (b2_x1 + b2_x2 - b1_x1 - b1_x2).pow(2)
        + (b2_y1 + b2_y2 - b1_y1 - b1_y2).pow(2)
    ) / 4

    v = (4 / math.pi ** 2) * ((w2 / h2).atan() - (w1 / h1).atan()).pow(2)
    with torch.no_grad():
        alpha = v / (1 - inner_iou + v + eps)

    return inner_iou - (rho2 / c2 + v * alpha)


def focal_iou_weight(iou, gamma=0.5):
    """Focal weight: 对高IoU样本加大权重。"""
    return iou.detach().pow(gamma)


def _inner_iou_forward(self, pred_dist, pred_bboxes, anchor_points,
                        target_bboxes, target_scores, target_scores_sum,
                        fg_mask, imgsz, stride):
    """替换 BboxLoss.forward，使用 Inner-IoU + Focal weighting。"""
    from ultralytics.utils.tal import bbox2dist

    weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)

    # Inner-IoU (ratio=0.7 对小目标效果好)
    iou = inner_bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask],
                          xywh=False, ratio=0.7)

    # Focal weighting: 高IoU样本获得更大权重
    focal_w = focal_iou_weight(iou.clamp(0, 1), gamma=0.5)
    loss_iou = ((1.0 - iou) * weight * focal_w).sum() / target_scores_sum

    # DFL loss (不变)
    if self.dfl_loss:
        target_ltrb = bbox2dist(anchor_points, target_bboxes, self.dfl_loss.reg_max - 1)
        loss_dfl = self.dfl_loss(pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max),
                                  target_ltrb[fg_mask]) * weight
        loss_dfl = loss_dfl.sum() / target_scores_sum
    else:
        target_ltrb = bbox2dist(anchor_points, target_bboxes)
        target_ltrb = target_ltrb * stride
        target_ltrb[..., 0::2] /= imgsz[1]
        target_ltrb[..., 1::2] /= imgsz[0]
        pred_dist_s = pred_dist * stride
        pred_dist_s[..., 0::2] /= imgsz[1]
        pred_dist_s[..., 1::2] /= imgsz[0]
        loss_dfl = (
            F.l1_loss(pred_dist_s[fg_mask], target_ltrb[fg_mask], reduction="none")
            .mean(-1, keepdim=True) * weight
        )
        loss_dfl = loss_dfl.sum() / target_scores_sum

    return loss_iou, loss_dfl


def patch_inner_iou():
    """Patch BboxLoss 使用 Inner-IoU + Focal。"""
    import ultralytics.utils.loss as los
    los.BboxLoss.forward = _inner_iou_forward
    print("[Loss] Patched BboxLoss → Inner-IoU (ratio=0.7) + Focal (gamma=0.5)")