NMS --classes 0 bug fix (#1710)
This commit is contained in:
1 parent
8bc0027afc
commit
1fc9d42a64
1 file changed
+1
-1
+1
-1
@@ -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]
|
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
||||||
|
|
||||||
# Filter by class
|
# Filter by class
|
||||||
if classes:
|
if classes is not None:
|
||||||
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
||||||
|
|
||||||
# Apply finite constraint
|
# Apply finite constraint
|
||||||
|
|||||||
Reference in new issue
Block a user