Confusion matrix native image-space fix (ultralytics#2046)

Make sure the labels and predictions are equally scaled on confusion_matrix.process_batch
taicaile · Jan 26, 2021 · 673fc9c · 673fc9c
1 parent 19d2707
commit 673fc9c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test.py b/test.py
@@ -178,7 +178,7 @@ def test(data,
                 tbox = xywh2xyxy(labels[:, 1:5])
                 scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1])  # native-space labels
                 if plots:
-                    confusion_matrix.process_batch(pred, torch.cat((labels[:, 0:1], tbox), 1))
+                    confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
 
                 # Per target class
                 for cls in torch.unique(tcls_tensor):