NMS --classes 0 bug fix (#1710)
Browse files- utils/general.py +1 -1
utils/general.py
CHANGED
@@ -311,7 +311,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
|
|
311 |
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
312 |
|
313 |
# Filter by class
|
314 |
-
if classes:
|
315 |
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
316 |
|
317 |
# Apply finite constraint
|
|
|
311 |
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
312 |
|
313 |
# Filter by class
|
314 |
+
if classes is not None:
|
315 |
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
316 |
|
317 |
# Apply finite constraint
|