import torch
def iou(box,boxes):
box_area = (box[2] - box[0]) * (box[3] - box[1])
bosex_area = (boxes[:,2] - boxes[:,0]) * (boxes[:,3] - boxes[:,1])
xx1 = torch.maximum(box[0],boxes[:,0])
xx2 = torch.minimum(box[2],boxes[:,2])
yy1 = torch.maximum(box[1],boxes[:,1])
yy2 = torch.minimum(box[3],boxes[:,3])
w,h = torch.maximum(torch.tensor(0),(xx2 - xx1)),torch.maximum(torch.tensor(0),(yy2 - yy1))
area = w * h
return area / (box_area + bosex_area - area)
def nms(boxes,thres):
new_boxes = boxes[boxes[:,5].argsort(descending=True)]
keep_box = []
while len(new_boxes) > 0:
_box = new_boxes[0]
keep_box.append(_box)
if len(new_boxes) > 1:
_boxes = new_boxes[1:]
new_boxes = _boxes[torch.where(iou(_box,_boxes) < thres)]
return torch.stack(keep_box)
if __name__ == "__main__":
box = torch.tensor([0,0,4,4])
boxes = torch.tensor([[4,4,7,7],[2,2,6,6]])
iou_ = iou(box,boxes)
NMS代码基础
最新推荐文章于 2024-08-14 22:29:49 发布