From c43be7b350cfff8f2423547c6aa0e8d6db07061b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 10 Sep 2018 15:58:01 +0200 Subject: [PATCH] updates --- test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test.py b/test.py index 1a9f8cb8..b55c039e 100644 --- a/test.py +++ b/test.py @@ -79,8 +79,11 @@ for batch_i, (imgs, targets) in enumerate(dataloader): # If no annotations add number of detections as incorrect if annotations.size(0) == 0: + target_cls = [] correct.extend([0 for _ in range(len(detections))]) else: + target_cls = annotations[:, 0] + # Extract target boxes as (x1, y1, x2, y2) target_boxes = xywh2xyxy(annotations[:, 1:5]) target_boxes *= opt.img_size @@ -101,7 +104,7 @@ for batch_i, (imgs, targets) in enumerate(dataloader): correct.append(0) # Compute Average Precision (AP) per class - target_cls = annotations[:, 0] if annotations.size(0) > 1 else annotations[0] + # target_cls = annotations[:, 0] if annotations.size(0) > 1 else annotations[0] AP = ap_per_class(tp=correct, conf=detections[:, 4], pred_cls=detections[:, 6], target_cls=target_cls) # Compute mean AP for this image