Skip to content

Commit 5f4ea0b

Browse files
authored
test.py native --single-cls (ultralytics#2928)
1 parent ee8b159 commit 5f4ea0b

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

test.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ def test(data,
119119
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
120120
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
121121
t = time_synchronized()
122-
out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=True)
122+
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
123123
t1 += time_synchronized() - t
124124

125125
# Statistics per image
@@ -136,6 +136,8 @@ def test(data,
136136
continue
137137

138138
# Predictions
139+
if single_cls:
140+
pred[:, 5] = 0
139141
predn = pred.clone()
140142
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
141143

0 commit comments

Comments
 (0)