本文主要是介绍YOLOv5 | NMS | 手动实现NMS，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
这个NMS可以按照相同的类别进行nms，不同类别是不会进行nms的，可以加一个“开关”（参数）来控制是否使用进行同类别iou，但是懒得加了哈哈哈哈 ，就是一个if语句，offset这个简直绝了，学的源码的哈哈哈哈哈哈，不信你去看
YOLOv5 | NMS | 源码解析-CSDN博客
import numpy as np
import torchdef nms(boxes, score, conf_thres, iou_thres, classes):tf = score[:] > conf_thresscore = score[tf]keep = []idx = np.argsort(score)[::-1]pred = torch.hstack((boxes,classes))offset = pred[:,4:5] * 4096boxes = pred[:,0:4] + offset# print(boxes)while len(idx) > 0:keep.append(idx[0])overlap = np.zeros_like(idx[1:], dtype= np.float32)for i, j in enumerate(idx[1:]):bbox1 = boxes[idx[0]]bbox2 = boxes[j]out = iou(bbox1, bbox2)overlap[i] = outidx = idx[1:][overlap < iou_thres]return keep# box = []# for i in keep:# box.append(boxes[i])# return boxdef iou(bbox1,bbox2):x1, y1, w1, h1 = bbox1 # left up right bottomx2, y2, w2, h2 = bbox2if (x1 <= ((x2+ w2)/2) <= w1) or ((y1 <= ((y2+ h2)/2) <= h1)): # 必须加上等号left_top_x = max(x1, x2)left_top_y = max(y1, y2) right_bottom_x = min(w1, w2)right_bottom_y = min(h1, h2)I = (right_bottom_x - left_top_x) * (right_bottom_y - left_top_y)o = (w1 - x1) * (h1- y1) + (w2 - x2) * (h2 - y2) - IIoU = I / oreturn IoUelse: # 不相交return 0if __name__ == '__main__':boxes = torch.Tensor([[100, 100, 200, 200],[120, 110, 220, 210],[300, 320, 400, 400],[180, 100, 300, 180]])scores = np.array([0.9, 0.8, 0.7, 0.6])classes = torch.Tensor([[0],[1],[1],[2]])# np.concatenate 先转置 可以拼接 out = nms(boxes, scores, 0.5, 0.25,classes)print(out)写的比较简单，不会高级的写法
这篇关于YOLOv5 | NMS | 手动实现NMS的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
