Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

test.py native --single-cls (#2928)

Browse Source
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-04-25 20:05:16 +02:00
1 parent aa78069c58
commit 3665c0f59b
1 file changed
+3 -1
Download Patch File Download Diff File Expand all files Collapse all files
test.py
+3 -1
View File
@@ -119,7 +119,7 @@ def test(data,
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels 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 lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_synchronized() 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 t1 += time_synchronized() - t
# Statistics per image # Statistics per image
@@ -136,6 +136,8 @@ def test(data,
continue continue
# Predictions # Predictions
if single_cls:
pred[:, 5] = 0
predn = pred.clone() predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred