yolov5 中build_targets函数的源码

以下是 YOLOv5 中 build_targets 函数的源码，该函数用于生成训练时所需的目标标签，其输入是一个包含了所有 Ground Truth 边界框的列表和一个预测的输出张量。它的输出是一个包含了每个预测边界框的目标类别、置信度得分以及边界框坐标调整量的张量。

def build_targets(pred, targets, hyp):
    # pred : (batch_size, num_anchors, grid_size, grid_size, 5 + num_classes)
    # targets : (num_targets, 6) [batch_index, class_id, x, y, w, h]
    # hyp : dict containing hyperparameters

    ignore_threshold = hyp['ignore_thresh']
    device = pred.device
    num_classes = hyp['num_classes']
    anchors = hyp['anchors']
    anchor_t = torch.tensor(anchors).float().to(device).view(1, -1, 1, 1, 2)

    b, a, gj, gi = targets[:, :4].long().t()
    target_cls = targets[:, 4].long()
    txywh = targets[:, 2:6]  # ground truth box
    gxy, gwh = txywh[:, :2], txywh[:, 2:]
    gij = (gxy * hyp['nw']).floor()

    # iou of targets-anchors (using wh only)
    box1 = txywh[:, None, :].repeat(1, anchor_t.shape[1], 1)
    box2 = anchor_t.repeat(gxy.shape[0], 1, 1, 1, 1).view(-1, 4)
    box2 = torch.cat((torch.zeros_like(box2[:, :2]), box2[:, 2:]), 1)  # convert [x,y,w,h] to [0,0,w,h]
    iou = bbox_iou(box1.view(-1, 4), box2, x1y1x2y2=False)  # iou(box, anchor)

    # anchor boxes with highest IoU
    topi, topk = iou.topk(1)

    # ensure each target is matched to highest iou
    a = topk % anchor_t.shape[1]  # anchor index
    b = topk // anchor_t.shape[1]  # batch index

    # XY coordinates
    gxy /= hyp['stride'][0]
    gi *= 0
    gj *= 0

    # target GT box XY coordinates with respect to cell
    txy = gxy - torch.cat((gxy.floor(),), 1)

    # Width and height (yolo method)
    tw, th = torch.sqrt(gwh / anchor_t[b, a, gj, gi]).unbind(1)
    # tw, th = torch.log(gwh / anchor_t[b, a, gj, gi]).unbind(1)  # yolo method

    # target GT box normalized width and height (yolo method)
    # twh = torch.log(gwh / anchor_t[b, a, gj, gi])
    twh = torch.sqrt(gwh / anchor_t[b, a, gj, gi])

    # cls, weight, class_mask
    tcls = torch.zeros_like(pred[..., 5:])
    tcls[b, a, gj, gi, target_cls] = 1.0
    tconf = torch.ones_like(pred[..., 4:5])
    tconf[b, a, gj, gi] = 0.0

    # iou of targets (using wh only)
    iou = bbox_iou(txywh, torch.cat((gij.float(), twh.float()), 1), x1y1x2y2=False)  # iou(target, anchor)

    # reject anchors below ignore_threshold iou
    tconf[iou < ignore_threshold] = 0.0

    # xy delta
    txy = txy - gij

    # wh delta
    twh = torch.log(gwh / anchor_t[b, a, gj, gi])
    twh[(iou < ignore_threshold).view(twh.shape)] = 0.0  # set non-matching boxes to 0

    # create output tensor
    output = torch.cat((txy, twh, tconf, tcls), -1)

    return output

该函数首先将预测的输出张量和目标 Ground Truth 边界框的信息进行解析和匹配，然后计算每个预测边界框与其对应的 Ground Truth 边界框的 IoU 值，最后根据阈值筛选出需要忽略的边界框以及需要进行训练的边界框，并生成相应的训练标签。具体实现可以参考以上源码注释。