瀏覽代碼

torch.ops.torchvision.nms (#860)

Don't load the entire torchvision library just for nms when the function is already in the torch library.
5.0
rafale77 GitHub 4 年之前
父節點
當前提交
9776e70988
沒有發現已知的金鑰在資料庫的簽署中 GPG Key ID: 4AEE18F83AFDEB23
共有 1 個文件被更改,包括 1 次插入2 次删除
  1. +1
    -2
      utils/general.py

+ 1
- 2
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)

Loading…
取消
儲存