Skip to content

Commit

Permalink
test.py native --single-cls (ultralytics#2928)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored Apr 25, 2021
1 parent a0c5fce commit 256472a
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion test.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def test(data,
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_synchronized()
out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=True)
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
t1 += time_synchronized() - t

# Statistics per image
Expand All @@ -136,6 +136,8 @@ def test(data,
continue

# Predictions
if single_cls:
pred[:, 5] = 0
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred

Expand Down

0 comments on commit 256472a

Please sign in to comment.