【Yolov5_Pytorch】相关修改

Yolov5相关修改

【1】修改NMS、添加DIOU_NMS、CIOU_NMS

general.py
其中的non_max_suppression函数中：

"""
其中boxes为Nx4的tensor，N为框的数量，4则为x1 y1 x2 y2
socres为N维的tensor，表示每个框的置信度
iou_thres则为上面算法中的IOU阈值
返回值为一个去除了过于相似框后的，根据置信度降序排列的列表，我们就可以根据此列表输出预测框
"""
i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS

添加NMS函数，后续可在此设置True或False启停：

def NMS(boxes, scores, iou_thres, GIoU=False, DIoU=False, CIoU=True):
    """
    :param boxes:  (Tensor[N, 4])): are expected to be in ``(x1, y1, x2, y2)
    :param scores: (Tensor[N]): scores for each one of the boxes
    :param iou_thres: discards all overlapping boxes with IoU > iou_threshold
    :return:keep (Tensor): int64 tensor with the indices
            of the elements that have been kept
            by NMS, sorted in decreasing order of scores
    """
    # Sort by conf from largest to smallest
    B = torch.argsort(scores, dim=-1, descending=True)
    keep = []
    while B.numel() > 0:
        # The highest confidence level is taken
        index = B[0]
        keep.append(index)
        if B.numel() == 1: break
        # Calculate iou,choose GIOU,DIOU,CIOU according to the demand
        iou = bbox_iou(boxes[index, :], boxes[B[1:], :], GIoU=GIoU, DIoU=DIoU, CIoU=CIoU)
        # Find the subscript that matches the threshold
        inds = torch.nonzero(iou <= iou_thres).reshape(-1)
        B = B[inds + 1]
    return torch.tensor(keep)

同时在non_max_suppression函数作如下修改：

i = NMS(boxes, scores, iou_thres)  # updata NMS
# i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS

参考：
（1）https://blog.csdn.net/lzzzzzzm/article/details/120151155
（2）https://blog.csdn.net/m0_53578855/article/details/124056523?utm_term=yolov5NMS%E6%8D%A2%E4%B8%BAWBF&utm_medium=distribute.pc_aggpage_search_result.none-task-blog-2_allsobaiduweb~default-0-124056523&spm=3001.4430

---

【2】分类置信度

general.py
添加函数：
其中 conf_thres1 列表为分别对应你的类别的置信度（此处是28分类的置信度列表）

def non_max_suppression_list(prediction, conf_thres1=[0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45,
                                                      0.45, 0.45, 0.45, 0.45, 0.45, 0.45,0.45, 0.45, 0.45, 0.45, 0.45, 0.45,
                                                      0.45, 0.45, 0.45, 0.45], iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
                        labels=()):
    """Runs Non-Maximum Suppression (NMS) on inference results

    Returns:
         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
    """

    nc = prediction.shape[2] - 5  # number of classes
    # xc = prediction[..., 4] > conf_thres  # candidates
    max_idx = torch.max(prediction[:, :, 5:], 2)[1]
    xc = prediction[..., 4] < 0
    for idx, conf_t in enumerate(conf_thres1):
        xc = xc | ((max_idx == idx) & (prediction[..., 4] > conf_t))

    # Settings
    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
    max_det = 300  # maximum number of detections per image
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 1.0  # seconds to quit after
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    merge = False  # use merge-NMS

    t = time.time()
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
        x = x[xc[xi]]  # confidence

        # If none remain process next image
        if not x.shape[0]:
            continue

        # Compute conf
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

        # Detections matrix nx6 (xyxy, conf, cls)
        if multi_label:
            # i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            i = torch.as_tensor(range(x.size(0)), device=x.device)
            j = torch.max(x[:, 5:], 1)[1]
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # best class only
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float()), 1)
            conf_temp = x[:, 4] < 0
            for idx, conf_t in enumerate(conf_thres1):
                conf_temp = conf_temp | ((x[:, 5] == idx) & (x[:, 4] > conf_t))
            x = x[conf_temp]

        # Check shape
        n = x.shape[0]  # number of boxes
        if not n:  # no boxes
            continue
        elif n > max_nms:  # excess boxes
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence

        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = NMS(boxes, scores, iou_thres)  # NMS
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]

        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            print(f'WARNING: NMS time limit {time_limit}s exceeded')
            break  # time limit exceeded

    return output

detect.py

from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, non_max_suppression_list

parser.add_argument('--conf-thres', type=float, default=[0.45, 0.8, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45, 0.45,
                                                      0.45, 0.45, 0.45, 0.45, 0.45, 0.45,0.45, 0.45, 0.45, 0.45, 0.45, 0.45,
                                                      0.45, 0.45, 0.45, 0.7], help='object confidence threshold')

参考：https://blog.csdn.net/shanglianlm/article/details/121034765

---

【3】未完