Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

torch.ops.torchvision.nms (#860)

Browse Source 
Don't load the entire torchvision library just for nms when the function is already in the torch library.
This commit is contained in:
rafale77 authored and GitHub committed 2020-08-27 11:13:41 -07:00
1 parent 0c01afc390
commit 9776e70988
1 file changed
+1 -2
Download Patch File Download Diff File Expand all files Collapse all files
utils/general.py
+1 -2
View File
@@ -17,7 +17,6 @@ import matplotlib.pyplot as plt
import numpy as np import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
import torchvision
import yaml import yaml
from scipy.cluster.vq import kmeans from scipy.cluster.vq import kmeans
from scipy.signal import butter, filtfilt 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 # Batched NMS
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores 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 if i.shape[0] > max_det: # limit detections
i = i[:max_det] i = i[:max_det]
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)