YOLOv3中do_nms_sort详解

总结起来就是三个过程:

1)去除没有物体的检测结果。

2)按照检测结果类别排序。

3)计算检测结果之间的iou,若大于thresh,则舍去。

参考:
1

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
YOLOv3 nms 函数用于过滤检测结果的重叠框,保留置信度最高的框。nms 函数的具体实现可以参考以下代码: ```python def non_max_suppression(prediction, confidence, num_classes, nms_conf=0.4): """ 使用非极大值抑制过滤检测结果的重叠框 """ # 获取所有检测框的左上角和右下角坐标 box_corner = prediction.new(prediction.shape) box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2 box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2 box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2 box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2 prediction[:, :, :4] = box_corner[:, :, :4] output = [None for _ in range(len(prediction))] # 对每个图片的检测结果进行处理 for image_i, image_pred in enumerate(prediction): # 获取该图片置信度最高的类别 max_conf, max_conf_idx = torch.max(image_pred[:, 5:5 + num_classes], 1) max_conf = max_conf.float().unsqueeze(1) max_conf_idx = max_conf_idx.float().unsqueeze(1) # 将置信度最高的类别和置信度分数合并到一起 seq = (image_pred[:, :5], max_conf, max_conf_idx) image_pred = torch.cat(seq, 1) # 删除置信度小于阈值的框 non_zero_indices = (torch.nonzero(image_pred[:, 4])) try: image_pred_ = image_pred[non_zero_indices.squeeze(), :].view(-1, 7) except: continue # 获取所有检测框的面积 img_classes = unique(image_pred_[:, -1]) for cls in img_classes: # 获取该类别的检测结果 cls_mask = image_pred_ * (image_pred_[:, -1] == cls).float().unsqueeze(1) class_mask_idx = torch.nonzero(cls_mask[:, -2]).squeeze() image_pred_class = image_pred_[class_mask_idx].view(-1, 7) # 按照置信度排序 conf_sort_index = torch.sort(image_pred_class[:, 4], descending=True)[1] image_pred_class = image_pred_class[conf_sort_index] idx = image_pred_class.size(0) # 检测框的数量 for i in range(idx): # 计算当前检测框与其它检测框的 IoU try: ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i + 1:]) except ValueError: break except IndexError: break # 将 IoU 大于阈值的检测框删除 iou_mask = (ious < nms_conf).float().unsqueeze(1) image_pred_class[i + 1:] *= iou_mask non_zero_indices = torch.nonzero(image_pred_class[:, 4]).squeeze() image_pred_class = image_pred_class[non_zero_indices].view(-1, 7) batch_index = image_pred_class.new(image_pred_class.size(0), 1).fill_(image_i) # 将图片索引加入到检测框 seq = batch_index, image_pred_class if not output[image_i]: output[image_i] = torch.cat(seq, 1) else: output[image_i] = torch.cat((output[image_i], torch.cat(seq, 1))) return output ``` 其,bbox_iou 函数用于计算两个框的 IoU 值,unique 函数用于获取一个张量的所有唯一值。使用时,可以将预测结果和对应的置信度输入到 non_max_suppression 函数进行处理,如下所示: ```python # 获取预测结果和置信度 prediction = model(x) conf_mask = (prediction[:, :, 4] > confidence_threshold).float().unsqueeze(2) prediction = prediction * conf_mask # 过滤掉重叠的检测框 output = non_max_suppression(prediction, confidence_threshold, num_classes) ```
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值