pytorch iou取最大

def iou(box1, box2):
    '''  
    box:[x1, y1, x2, y2]
    '''
    in_h = min(box1[3], box2[3]) - max(box1[1], box2[1])
    in_w = min(box1[2], box2[2]) - max(box1[0], box2[0])
    inner = 0 if in_h<0 or in_w<0 else in_h*in_w
    union = (box1[2] - box1[0]) * (box1[3] - box1[1]) + \
            (box2[2] - box2[0]) * (box2[3] - box2[1]) - inner
    print(inner)
    iou = inner / union
    return iou
box1 = [0,0,10,10]
box2 = [5,5,15,15]
print(iou(box1,box2))

import torch

boxs=torch.Tensor([[0.1,0.5,0.3],[0.2,0.3,0.4],[0.5,0.6,0.7]])# 3个框

iou=torch.Tensor([[0.1,0.5,0.3,0.2],[0.2,0.4,0.6,0.1],[0.2,0.1,0.4,0.1]])


print(boxs.size())
b1,b2=iou.max(0)
print(b1)
print(b2)
c=boxs[b2]
print(c)