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

--classes bug fix #17

Browse Source
This commit is contained in:
Glenn Jocher committed 2020-06-05 13:29:26 -07:00
1 parent 22d6088205
commit 948bcdd219
1 file changed
+2 -1
Download Patch File Download Diff File Expand all files Collapse all files
utils/utils.py
+2 -1
View File
@@ -460,6 +460,7 @@ def build_targets(p, targets, model):
return tcls, tbox, indices, anch return tcls, tbox, indices, anch
def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, fast=False, classes=None, agnostic=False): def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, fast=False, classes=None, agnostic=False):
""" """
Performs Non-Maximum Suppression on inference results Performs Non-Maximum Suppression on inference results
@@ -508,7 +509,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, fast=False, c
# Filter by class # Filter by class
if classes: if classes:
x = x[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)] x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
# Apply finite constraint # Apply finite constraint
# if not torch.isfinite(x).all(): # if not torch.isfinite(x).all():