From 948bcdd21918043f309d763610d51a91afe970d2 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 5 Jun 2020 13:29:26 -0700 Subject: [PATCH] --classes bug fix #17 --- utils/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/utils.py b/utils/utils.py index e22cf14a4973..659c939b4c6a 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -460,6 +460,7 @@ def build_targets(p, targets, model): return tcls, tbox, indices, anch + 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 @@ -508,7 +509,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, fast=False, c # Filter by class 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 # if not torch.isfinite(x).all():