Browse Source

NMS --classes 0 bug fix (#1710)

5.0
Polydefkis Gkagkos GitHub 3 years ago
parent
commit
1fc9d42a64
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      utils/general.py

+ 1
- 1
utils/general.py View File

@@ -311,7 +311,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]

# Filter by class
if classes:
if classes is not None:
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]

# Apply finite constraint

Loading…
Cancel
Save