test.py native --single-cls (#2928)
Browse files
test.py
CHANGED
@@ -119,7 +119,7 @@ def test(data,
|
|
119 |
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
120 |
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
121 |
t = time_synchronized()
|
122 |
-
out = non_max_suppression(out, conf_thres
|
123 |
t1 += time_synchronized() - t
|
124 |
|
125 |
# Statistics per image
|
@@ -136,6 +136,8 @@ def test(data,
|
|
136 |
continue
|
137 |
|
138 |
# Predictions
|
|
|
|
|
139 |
predn = pred.clone()
|
140 |
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
|
141 |
|
|
|
119 |
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
120 |
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
121 |
t = time_synchronized()
|
122 |
+
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
|
123 |
t1 += time_synchronized() - t
|
124 |
|
125 |
# Statistics per image
|
|
|
136 |
continue
|
137 |
|
138 |
# Predictions
|
139 |
+
if single_cls:
|
140 |
+
pred[:, 5] = 0
|
141 |
predn = pred.clone()
|
142 |
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
|
143 |
|