yolov5改进softnms

### YOLOv5 中改进 SoftNMS 算法实现方法 在YOLOv5中，为了提升检测精度并减少冗余边界框的数量，可以采用Soft-NMS来替代传统的非极大值抑制(NMS)[^1]。具体来说，在传统NMS基础上引入了更平滑的方式调整候选框的置信度分数。 对于每一个预测框\(b_i\)以及当前得分最高的框\(M\)，如果两者之间的交并比(IOU)较大，则适当减小该预测框的置信度分值而不是直接丢弃它。这种做法能够更好地保留一些重叠但仍然可能是不同物体实例的边界框[^2]。 以下是Python代码片段展示了如何修改YOLOv5中的`non_max_suppression()`函数以支持Soft NMS： ```python def soft_non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, method='linear'): """ prediction: (tensor), shape(batch_size, num_boxes, 5+num_classes) box format is center_x, center_y, width, height, object_confidence. Returns detections with shape: (batch_index, class_index, x_center, y_center, w, h, obj_conf, class_conf) Method can be 'gaussian' or 'linear'. """ output = [None] * prediction.shape[0] for image_idx, pred in enumerate(prediction): xc = pred[..., 4] > conf_thres # candidates pred = pred[xc] max_wh = 4096 # maximum filter wh size if not pred.size(0): continue boxes = xywh2xyxy(pred[:, :4]) # center_x,center_y,w,h to xmin,ymin,xmax,ymax scores = pred[:, 4].sigmoid() * pred[:, 5:].softmax(dim=-1).max(-1)[0] classes = pred[:, 5:].argmax(axis=-1) indices = torchvision.ops.batched_nms( boxes, scores, classes, iou_threshold=iou_thres ) dets = torch.cat((pred[indices], scores[indices, None]), dim=1) keep = [] while len(dets) > 0: large_overlap = bbox_iou(dets[0][None, :4], dets[:, :4]).view(-1) >= iou_thres current_box = dets[0] if method == "linear": weights = (1-large_overlap.view(-1)).clamp(min=0.) elif method == "gaussian": sigma = 0.5 weights = np.exp(-(bbox_iou(current_box[:4][None,:],dets[:,:4])**2)/sigma**2) dets[large_overlap, 4:] *= weights.unsqueeze(-1) keep.append(indices[torch.argmax(scores)]) valid_scores_mask = dets[:, 4]>conf_thres dets=dets[valid_scores_mask] indices=indices[valid_scores_mask] scores=scores[valid_scores_mask] output[image_idx] = pred[keep] return output ``` 此版本实现了两种不同的权重衰减方式："线性"(Linear)和"高斯"(Gaussian)，可以根据实际需求选择合适的方法应用于目标检测任务中。