From 9776e7098880903e0549376a7b62cfdc66b4ec12 Mon Sep 17 00:00:00 2001 From: rafale77 Date: Thu, 27 Aug 2020 11:13:41 -0700 Subject: [PATCH] torch.ops.torchvision.nms (#860) Don't load the entire torchvision library just for nms when the function is already in the torch library. --- utils/general.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/utils/general.py b/utils/general.py index 79b843e..36d52e3 100755 --- a/utils/general.py +++ b/utils/general.py @@ -17,7 +17,6 @@ import matplotlib.pyplot as plt import numpy as np import torch import torch.nn as nn -import torchvision import yaml from scipy.cluster.vq import kmeans from scipy.signal import butter, filtfilt @@ -651,7 +650,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, # Batched NMS c = x[:, 5:6] * (0 if agnostic else max_wh) # classes boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores - i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) + i = torch.ops.torchvision.nms(boxes, scores, iou_thres) if i.shape[0] > max_det: # limit detections i = i[:max_det] if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)