| """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 |
|
|
| |
| 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_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_w1 * inner_h1 + inner_w2 * inner_h2 - inner_inter + eps |
|
|
| |
| inner_iou = inner_inter / inner_union |
|
|
| |
| 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) |
|
|
| |
| iou = inner_bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], |
| xywh=False, ratio=0.7) |
|
|
| |
| focal_w = focal_iou_weight(iou.clamp(0, 1), gamma=0.5) |
| loss_iou = ((1.0 - iou) * weight * focal_w).sum() / target_scores_sum |
|
|
| |
| 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)") |
|
|
