torchvision nms

要求格式[x1,y1,x2,y2]

if __name__ == '__main__':

    import torch
    a=torch.Tensor([[1,1,2,2],[1,1,3.100001,3],[1,1,3.1,3]])
    b=torch.Tensor([[0.9],[0.98],[0.980005]])

    from torchvision.ops import nms

    ccc=nms(a,b,0.4)
    print(ccc)
    print(a[ccc])

 

from __future__ import print_function
import os
import argparse
import torch
import torch.backends.cudnn as cudnn
import numpy as np
from torchvision.ops import nms

from data import cfg_mnet, cfg_re50, cfg_peleenet
from layers.functions.prior_box import PriorBox
from utils.nms.py_cpu_nms import py_cpu_nms
import cv2
from models.retinaface import RetinaFace
from utils.box_utils import decode, decode_landm
import time

parser = argparse.ArgumentParser(description='Retinaface')

# parser.add_argument('-m', '--trained_model', default='weights/mobilenet0.25_Final.p