YOLOv5对CIOU损失函数改进的代码

以下是YOLOv5对CIOU损失函数改进的代码，供您参考：

import torch.nn.functional as F
from utils.general import box_iou

def compute_ciou_loss(pred, gt, eps=1e-7):
    """
    Args：
        pred: 预测的边界框，shape为[N, 4], [x, y, w, h]
        gt: 真实的边界框，shape为[N, 4], [x, y, w, h]
    """
    # 将边界框的坐标转换为左上角和右下角的点的坐标
    pred_xy = pred[:, :2]
    pred_wh = pred[:, 2:]
    pred_mins = pred_xy - pred_wh / 2.0
    pred_maxs = pred_xy + pred_wh / 2.0

    gt_xy = gt[:, :2]
    gt_wh = gt[:, 2:]
    gt_mins = gt_xy - gt_wh / 2.0
    gt_maxs = gt_xy + gt_wh / 2.0

    # 计算真实边界框和预测边界框的交集
    intersect_mins = torch.max(pred_mins, gt_mins)
    intersect_maxs = torch.min(pred_maxs, gt_maxs)
    intersect_wh = torch.clamp(intersect_maxs - intersect_mins, min=0)
    intersect_area = intersect_wh[:, 0] * intersect_wh[:, 1]

    # 计算真实边界框和预测边界框的并集
    pred_area = pred_wh[:, 0] * pred_wh[:, 1]
    gt_area = gt_wh[:, 0] * gt_wh[:, 1]
    union_area = pred_area + gt_area - intersect_area

    # 计算IoU
    iou = intersect_area / (union_area + eps)

    # 计算中心点的距离
    center_distance = torch.sum(torch.pow((pred_xy - gt_xy), 2), axis=1)

    # 计算最小外接矩形的对角线长度
    enclose_mins = torch.min(pred_mins, gt_mins)
    enclose_maxs = torch.max(pred_maxs, gt_maxs)
    enclose_wh = torch.clamp(enclose_maxs - enclose_mins, min=0)
    enclose_diagonal = torch.sum(torch.pow(enclose_wh, 2), axis=1)

    # 计算CIOU
    ciou = iou - center_distance / enclose_diagonal

    # 计算损失
    ciou_loss = 1 - ciou
    return ciou_loss

在YOLOv5的损失函数中，我们可以使用该函数来计算CIOU损失。