Unverified 提交 1fc9d42a authored 作者: Polydefkis Gkagkos's avatar Polydefkis Gkagkos 提交者: GitHub

NMS --classes 0 bug fix (#1710)

上级 8bc0027a
...@@ -311,7 +311,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non ...@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论